79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
from typing import Any, Callable, Dict, Generic, List, NewType, Tuple, Type, TypeVar, Union
|
|
|
|
from pydantic import BaseModel
|
|
from tortoise.expressions import Q
|
|
from tortoise.models import Model
|
|
|
|
Total = NewType("Total", int)
|
|
ModelType = TypeVar("ModelType", bound=Model)
|
|
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
|
|
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)
|
|
|
|
|
|
class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
|
|
def __init__(self, model: Type[ModelType]):
|
|
self.model = model
|
|
|
|
async def is_exist(self, **kwargs) -> bool:
|
|
return await self.model.filter(**kwargs).first()
|
|
|
|
async def all(self, search: Q = Q()) -> List[ModelType]:
|
|
return await self.model.filter(search).all()
|
|
|
|
async def get(self, id: int) -> ModelType:
|
|
return await self.model.get(id=id)
|
|
|
|
async def list(self, page: int, page_size: int, search: Q = Q(), order: list = []) -> Tuple[Total, List[ModelType]]:
|
|
query = self.model.filter(search)
|
|
return await query.count(), await query.offset((page - 1) * page_size).limit(page_size).order_by(*order)
|
|
|
|
async def create(self, obj_in: CreateSchemaType) -> ModelType:
|
|
if isinstance(obj_in, Dict):
|
|
obj_dict = obj_in
|
|
else:
|
|
obj_dict = obj_in.model_dump()
|
|
obj = self.model(**obj_dict)
|
|
await obj.save()
|
|
return obj
|
|
|
|
async def copy(self, id: int, handler: Callable[[Dict[str, Any]], Dict[str, Any]] = None) -> ModelType:
|
|
obj = await self.get(id=id)
|
|
obj_dict = await obj.to_dict()
|
|
drop_keys = ["id", "created_at", "updated_at"]
|
|
for key in drop_keys:
|
|
obj_dict.pop(key, None)
|
|
if handler:
|
|
obj_dict = handler(obj_dict)
|
|
obj = self.model(**obj_dict)
|
|
await obj.save()
|
|
return obj
|
|
|
|
async def update(self, id: int, obj_in: Union[UpdateSchemaType, Dict[str, Any]]) -> ModelType:
|
|
if isinstance(obj_in, Dict):
|
|
obj_dict = obj_in
|
|
else:
|
|
obj_dict = obj_in.model_dump(exclude_unset=True, exclude={"id"})
|
|
obj = await self.get(id=id)
|
|
obj = obj.update_from_dict(obj_dict)
|
|
await obj.save()
|
|
return obj
|
|
|
|
async def create_or_update(self, obj_in: CreateSchemaType, query_kwargs: Dict[str, Any], update_kwargs: Dict[str, Any] = None) -> ModelType:
|
|
update_kwargs = update_kwargs or {}
|
|
orm = await self.model.filter(**query_kwargs).first()
|
|
if orm:
|
|
need_update = False if update_kwargs else True
|
|
for key, value in update_kwargs.items():
|
|
if getattr(orm, key) != value:
|
|
need_update = True
|
|
break
|
|
if need_update:
|
|
return 'update', await self.update(orm.id, obj_in)
|
|
else:
|
|
return False, orm
|
|
return 'create', await self.create(obj_in)
|
|
|
|
async def remove(self, id: int) -> None:
|
|
obj = await self.get(id=id)
|
|
await obj.delete()
|