from collections.abc import Sequence from datetime import datetime, timedelta from typing import Any, TypeVar, cast from pydantic import BaseModel from sqlalchemy import asc, delete, desc, false, func, literal_column, select, update from sqlalchemy import inspect as sa_inspect from sqlalchemy.engine import Result from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import load_only, selectinload from sqlalchemy.sql.elements import ColumnElement from app.core.base_model import ModelMixin from app.core.base_schema import AuthSchema, PageResultSchema from app.core.exceptions import CustomException OutSchemaType = TypeVar("OutSchemaType", bound=BaseModel) CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel) UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel) class CRUDBase[ModelType: ModelMixin, CreateSchemaType, UpdateSchemaType]: """统一数据层基类 核心设计:``auth`` 是必填的 ``AuthSchema``,子类可直接访问 ``self.auth.user.xxx``。 用法: class UserCRUD(CRUDBase[UserModel, UserCreateSchema, UserUpdateSchema]): def __init__(self, auth: AuthSchema, db: AsyncSession) -> None: super().__init__(model=UserModel, auth=auth, db=db) """ def __init__(self, model: type[ModelType], auth: AuthSchema, db: AsyncSession) -> None: self.model = model self.auth = auth self.db = db def _get_pk_col(self) -> ColumnElement: mapper = sa_inspect(self.model) pk_cols = list[Any](getattr(mapper, "primary_key", [])) if not pk_cols: raise CustomException(msg="模型缺少主键") if len(pk_cols) > 1: raise CustomException(msg="暂不支持复合主键操作") return pk_cols[0] @property def _supports_soft_delete(self) -> bool: return all(hasattr(self.model, attr) for attr in ("is_deleted", "deleted_time", "deleted_id")) def _soft_delete_values(self) -> dict[str, Any]: data: dict[str, Any] = {"is_deleted": True, "deleted_time": datetime.now()} if self.auth.user.id: data["deleted_id"] = self.auth.user.id return data async def get(self, preload: list[str | Any] | None = None, **kwargs) -> ModelType | None: try: conditions = await self.__build_conditions(**kwargs) sql = select(self.model).where(*conditions) for opt in self.__loader_options(preload): sql = sql.options(opt) result: Result = await self.db.execute(sql) return result.scalars().first() except Exception as e: raise CustomException(msg=f"获取查询失败: {e!s}") async def get_by_id(self, model_id: int) -> ModelType | None: return await self.get(id=model_id) async def get_or_404( self, id: int | None = None, msg: str = "该数据不存在", preload: list[str | Any] | None = None, out_schema: type[OutSchemaType] | None = None, **kwargs, ) -> ModelType | OutSchemaType: if id is not None: kwargs["id"] = id obj = await self.get(preload=preload, **kwargs) if not obj: raise CustomException(msg=msg) return out_schema.model_validate(obj) if out_schema else obj async def exists(self, **kwargs) -> bool: return await self.get(**kwargs) is not None async def count(self, **kwargs) -> int: try: conditions = await self.__build_conditions(**kwargs) count_sql = select(func.count()).select_from(self.model).where(*conditions) result: Result = await self.db.execute(count_sql) return result.scalar() or 0 except Exception as e: raise CustomException(msg=f"统计失败: {e!s}") async def get_list( self, search: dict[str, Any] | None = None, order_by: list[dict[str, str]] | None = None, preload: list[str | Any] | None = None, load_columns: list | None = None, ) -> Sequence[ModelType]: try: conditions = await self.__build_conditions(**(search or {})) order = order_by or [{"id": "asc"}] sql = select(self.model).where(*conditions).order_by(*self._parse_order(order)) if load_columns: sql = sql.options(load_only(*load_columns)) for opt in self.__loader_options(preload): sql = sql.options(opt) result: Result = await self.db.execute(sql) return result.scalars().all() except Exception as e: raise CustomException(msg=f"列表查询失败: {e!s}") async def tree_list( self, search: dict[str, Any] | None = None, order_by: list[dict[str, str]] | None = None, children_attr: str | None = None, preload: list[str | Any] | None = None, ) -> Sequence[ModelType]: if children_attr is None: children_attr = getattr(self.model, "__tree_children_attr__", "children") try: conditions = await self.__build_conditions(**(search or {})) order = order_by or [{"id": "asc"}] sql = select(self.model).where(*conditions).order_by(*self._parse_order(order)) final_preload = preload if preload is None and children_attr and hasattr(self.model, children_attr): model_defaults = getattr(self.model, "__loader_options__", []) final_preload = [*list(model_defaults), children_attr] for opt in self.__loader_options(final_preload): sql = sql.options(opt) result: Result = await self.db.execute(sql) return result.scalars().all() except Exception as e: raise CustomException(msg=f"树形列表查询失败: {e!s}") async def page( self, offset: int, limit: int, order_by: list[dict[str, str]], search: dict[str, Any] | None = None, out_schema: type[OutSchemaType] | None = None, preload: list[str | Any] | None = None, load_columns: list | None = None, ) -> PageResultSchema[OutSchemaType] | PageResultSchema: try: conditions = await self.__build_conditions(**(search or {})) order = order_by or [{"id": "asc"}] mapper = sa_inspect(self.model) pk_cols = list(getattr(mapper, "primary_key", [])) pk = pk_cols[0] if pk_cols else literal_column("1") data_sql = select(self.model).where(*conditions) if load_columns: data_sql = data_sql.options(load_only(*load_columns)) for opt in self.__loader_options(preload): data_sql = data_sql.options(opt) count_sql = select(func.count(pk)).select_from(self.model) where_clause = data_sql.whereclause if where_clause is not None: count_sql = count_sql.where(where_clause) total_result = await self.db.execute(count_sql) total = total_result.scalar() or 0 result: Result = await self.db.execute(data_sql.order_by(*self._parse_order(order)).offset(offset).limit(limit)) objs = result.scalars().all() items = [out_schema.model_validate(obj) for obj in objs] if out_schema else list(objs) return PageResultSchema( page_no=offset // limit + 1 if limit else 1, page_size=limit or 10, total=total, has_next=offset + limit < total, items=items, ) except Exception as e: raise CustomException(msg=f"分页查询失败: {e!s}") async def create(self, data: CreateSchemaType) -> ModelType: try: obj_dict = data if isinstance(data, dict) else cast("BaseModel", data).model_dump() obj = self.model(**obj_dict) user = self.auth.user if user.id: if hasattr(obj, "created_id"): setattr(obj, "created_id", user.id) if hasattr(obj, "updated_id"): setattr(obj, "updated_id", user.id) self.db.add(obj) await self.db.flush() await self.db.refresh(obj) return obj except Exception as e: raise CustomException(msg=f"创建失败: {e!s}") async def update(self, id: int, data: UpdateSchemaType) -> ModelType: try: obj_dict = data if isinstance(data, dict) else cast("BaseModel", data).model_dump(exclude_unset=True, exclude={"id"}) model_defaults = getattr(self.model, "__loader_options__", []) obj = await self.get(id=id, preload=model_defaults) if not obj: raise CustomException(msg="更新对象不存在") user = self.auth.user if user.id and hasattr(obj, "updated_id"): setattr(obj, "updated_id", user.id) for key, value in obj_dict.items(): if hasattr(obj, key): setattr(obj, key, value) await self.db.flush() await self.db.refresh(obj) return obj except CustomException: raise except Exception as e: raise CustomException(msg=f"更新失败: {e!s}") async def delete(self, ids: list[int]) -> None: try: pk = self._get_pk_col() if self._supports_soft_delete: sql = update(self.model).where(pk.in_(ids)).values(**self._soft_delete_values()) else: sql = delete(self.model).where(pk.in_(ids)) await self.db.execute(sql) await self.db.flush() except Exception as e: raise CustomException(msg=f"删除失败: {e!s}") async def clear(self) -> None: try: if self._supports_soft_delete: sql = update(self.model).values(**self._soft_delete_values()) else: sql = delete(self.model) await self.db.execute(sql) await self.db.flush() except Exception as e: raise CustomException(msg=f"清空失败: {e!s}") async def set(self, ids: list[int], **kwargs) -> None: try: pk = self._get_pk_col() sql = update(self.model).where(pk.in_(ids)).values(**kwargs) await self.db.execute(sql) await self.db.flush() except Exception as e: raise CustomException(msg=f"批量更新失败: {e!s}") async def restore(self, ids: list[int]) -> None: try: if not self._supports_soft_delete: raise CustomException(msg="该模型不支持软删除,无法恢复") pk = self._get_pk_col() sql = update(self.model).where(pk.in_(ids)).values(is_deleted=False, deleted_time=None, deleted_id=None) await self.db.execute(sql) await self.db.flush() except Exception as e: raise CustomException(msg=f"恢复失败: {e!s}") async def __build_conditions(self, **kwargs) -> list[ColumnElement]: conditions: list[ColumnElement] = [] if hasattr(self.model, "is_deleted"): conditions.append(getattr(self.model, "is_deleted") == false()) for key, value in kwargs.items(): if value is None or value == "": continue attr = getattr(self.model, key) if isinstance(value, tuple): conditions.extend(self._resolve_condition(attr, value)) elif isinstance(value, str): conditions.append(attr.like(f"%{value}%")) elif isinstance(value, (int, bool)): conditions.append(attr == value) elif isinstance(value, list) and len(value) == 2: conditions.append(attr.between(value[0], value[1])) else: conditions.append(attr == value) return conditions @staticmethod def _resolve_condition(attr: ColumnElement, value: tuple) -> list[ColumnElement]: seq, val = value handlers: dict[str, Any] = { "None": lambda: [attr.is_(None)], "not None": lambda: [attr.isnot(None)], } if seq in handlers: return handlers[seq]() if val is None: return [] if seq == "date": dt = datetime.strptime(val, "%Y-%m-%d") return [attr >= dt, attr < dt + timedelta(days=1)] if seq == "month": dt = datetime.strptime(val, "%Y-%m") next_month = dt.replace(year=dt.year + 1, month=1) if dt.month == 12 else dt.replace(month=dt.month + 1) return [attr >= dt, attr < next_month] if seq == "like": return [attr.like(f"%{val}%")] if seq == "in": if isinstance(val, (list, tuple, set)) and len(val) == 0: return [false()] return [attr.in_(val)] if seq == "between" and isinstance(val, (list, tuple)) and len(val) == 2: return [attr.between(val[0], val[1])] _COMPARATORS = { "!=": attr.__ne__, "ne": attr.__ne__, ">": attr.__gt__, "gt": attr.__gt__, ">=": attr.__ge__, "ge": attr.__ge__, "<": attr.__lt__, "lt": attr.__lt__, "<=": attr.__le__, "le": attr.__le__, "eq": attr.__eq__, "==": attr.__eq__, } cmp = _COMPARATORS.get(seq) if cmp is not None: return [cmp(val)] return [] def _parse_order(self, order: list[dict[str, str]]) -> list[ColumnElement]: columns: list[ColumnElement] = [] for item in order: for field, direction in item.items(): column = getattr(self.model, field) # type: ignore[arg-type] columns.append(desc(column) if direction.lower() == "desc" else asc(column)) return columns def __loader_options(self, preload: list[str | Any] | None = None) -> list[Any]: model_loader_options = getattr(self.model, "__loader_options__", []) if preload == []: return [] names: set[str] = set(model_loader_options) if preload: for opt in preload: if isinstance(opt, str): names.add(opt) options: list[Any] = [] for name in names: if hasattr(self.model, name): options.append(selectinload(getattr(self.model, name))) if preload: options.extend(opt for opt in preload if not isinstance(opt, str)) return options