FastapiAdmin/backend/app/core/base_crud.py

from collections.abc import Sequence
from datetime import datetime, timedelta
from typing import Any, TypeVar, cast

from pydantic import BaseModel
from sqlalchemy import Select, 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
from app.core.permission import Permission

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:
        """初始化 CRUDBase。

        参数:
        - model: 数据模型类
        - auth: 认证信息
        - db: 数据库会话
        """
        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_one(self, preload: list[str | Any] | None = None, **kwargs) -> ModelType | None:
        """内部方法：在当前实例会话上执行单条查询（get / update 共用）

        参数:
        - preload: 预加载关系
        - **kwargs: 查询条件

        返回:
        - 对象实例或 None
        """
        conditions = await self.__build_conditions(**kwargs)
        sql = select(self.model).where(*conditions)
        for opt in self.__loader_options(preload):
            sql = sql.options(opt)
        sql = await self.__filter_permissions(sql)
        result: Result = await self.db.execute(sql)
        return result.scalars().first()

    async def get(self, preload: list[str | Any] | None = None, **kwargs) -> ModelType | None:
        """根据条件获取单个对象（复用请求级事务会话，保证读已写一致性）

        参数:
        - preload: 预加载关系
        - **kwargs: 查询条件

        返回:
        - 对象实例或 None
        """
        try:
            return await self._get_one(preload=preload, **kwargs)
        except CustomException:
            raise
        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:
        """按条件查询单条记录，不存在时抛出 404。

        参数:
        - id: 主键 ID（快捷方式，等价于 kwargs={"id": id}）。
        - msg: 不存在时的错误消息。
        - preload: 预加载关系列表。
        - out_schema: 输出 Schema，为 None 时返回 ORM 对象。
        - **kwargs: 其他查询条件（与 id 互斥）。

        返回:
        - ORM 对象或 Pydantic Schema 实例。

        异常:
        - CustomException: 记录不存在。
        """
        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:
        """检查是否存在符合条件的记录

        参数:
        - **kwargs: 查询条件

        返回:
        - 是否存在
        """
        return await self.get(**kwargs) is not None

    async def count(self, **kwargs) -> int:
        """统计符合条件的记录数（复用请求级事务会话）

        参数:
        - **kwargs: 查询条件，支持元组语法

        返回:
        - 记录数
        """
        try:
            conditions = await self.__build_conditions(**kwargs)
            count_sql = select(func.count()).select_from(self.model).where(*conditions)
            count_sql = await self.__filter_permissions(count_sql)
            result: Result = await self.db.execute(count_sql)
            return result.scalar() or 0
        except CustomException:
            raise
        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]:
        """根据条件获取对象列表（复用请求级事务会话）

        参数:
        - search: 查询条件
        - order_by: 排序字段, 格式为 [{'id': 'asc'}, {'name': 'desc'}]
        - preload: 预加载关系
        - load_columns: 仅加载指定的列（减少 SELECT 传输量）

        返回:
        - 对象列表
        """
        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)
            sql = await self.__filter_permissions(sql)
            result: Result = await self.db.execute(sql)
            return result.scalars().all()
        except CustomException:
            raise
        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]:
        """获取树形结构数据列表（复用请求级事务会话）

        参数:
        - search: 查询条件
        - order_by: 排序字段
        - children_attr: 子节点属性名（None 时自动从模型 __tree_children_attr__ 推断）
        - preload: 额外预加载关系

        返回:
        - 树形结构数据列表
        """
        # 自动从模型推断 children_attr
        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)

            sql = await self.__filter_permissions(sql)
            result: Result = await self.db.execute(sql)
            return result.scalars().all()
        except CustomException:
            raise
        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:
        """获取分页数据（复用请求级事务会话；count 与 data 共享同一会话）

        参数:
        - offset: 偏移量
        - limit: 每页数量
        - order_by: 排序字段
        - search: 查询条件
        - out_schema: 输出数据模型（None 时返回原始 ORM 对象）
        - preload: 预加载关系
        - load_columns: 仅加载指定的列（减少 SELECT 传输量）

        返回:
        - 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)
            data_sql = await self.__filter_permissions(data_sql)

            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 CustomException:
            raise
        except Exception as e:
            raise CustomException(msg=f"分页查询失败: {e!s}")

    async def create(self, data: CreateSchemaType) -> ModelType:
        """创建新对象（有认证时自动填充租户与审计字段）

        事务由 request 级 db_getter 统一管理，本方法不开启独立事务。

        参数:
        - data: 对象属性

        返回:
        - 新创建的对象实例
        """
        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, "tenant_id"):
                    # 仅当调用方未显式指定 tenant_id 时，才默认使用当前用户的租户
                    # 超管可以显式传任意 tenant_id（管理跨租户数据），非超管必须强制为本租户
                    if not hasattr(obj, "tenant_id") or getattr(obj, "tenant_id", None) is None:
                        setattr(obj, "tenant_id", user.tenant_id)
                    elif not user.is_superuser and getattr(obj, "tenant_id") != user.tenant_id:
                        raise CustomException(msg="无权创建其他租户的数据")
                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 CustomException:
            raise
        except Exception as e:
            raise CustomException(msg=f"创建失败: {e!s}")

    async def update(self, id: int, data: UpdateSchemaType) -> ModelType:
        """更新对象（有认证时检查租户归属 + 填充审计字段）

        事务由 request 级 db_getter 统一管理，本方法不开启独立事务。

        参数:
        - id: 对象 ID
        - data: 更新属性

        返回:
        - 更新后的对象实例
        """
        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_one(id=id, preload=model_defaults)
            if not obj:
                raise CustomException(msg="更新对象不存在")

            # 租户权限检查（仅在有认证且非超管时）
            user = self.auth.user
            if user.id and not user.is_superuser:
                if hasattr(obj, "tenant_id"):
                    obj_tid = getattr(obj, "tenant_id", None)
                    if obj_tid is not None and obj_tid != user.tenant_id:
                        is_platform = getattr(self.model, "__platform_data_shared__", False)
                        if is_platform and obj_tid == 1:
                            raise CustomException(msg="平台数据仅管理员可修改")
                        raise CustomException(msg="无权修改其他租户的数据")

            # 审计字段
            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 = self._tenant_dml_where(update(self.model).where(pk.in_(ids))).values(**self._soft_delete_values())
                await self.db.execute(sql)
            else:
                sql = self._tenant_dml_where(delete(self.model).where(pk.in_(ids)))
                await self.db.execute(sql)
            await self.db.flush()
        except CustomException:
            raise
        except Exception as e:
            raise CustomException(msg=f"删除失败: {e!s}")

    async def clear(self) -> None:
        """软清空对象表（有认证时填充删除人 + 租户隔离）"""
        try:
            if self._supports_soft_delete:
                sql = self._tenant_dml_where(update(self.model)).values(**self._soft_delete_values())
                await self.db.execute(sql)
            else:
                sql = self._tenant_dml_where(delete(self.model))
                await self.db.execute(sql)
            await self.db.flush()
        except CustomException:
            raise
        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 = self._tenant_dml_where(update(self.model)).where(pk.in_(ids)).values(**kwargs)
            await self.db.execute(sql)
            await self.db.flush()
        except CustomException:
            raise
        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 = self._tenant_dml_where(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 CustomException:
            raise
        except Exception as e:
            raise CustomException(msg=f"恢复失败: {e!s}")

    async def __filter_permissions(self, sql: Select) -> Select:
        """过滤数据权限（仅用于 Select）"""
        if not self.auth:
            return sql
        if getattr(self.model, "__platform_data_shared__", False):
            for condition in self._platform_shared_conditions():
                sql = sql.where(condition)
        filter_obj = Permission(model=self.model, auth=self.auth, db=self.db)
        return await filter_obj.filter_query(sql)

    def _platform_shared_conditions(self) -> list[ColumnElement]:
        user = self.auth.user
        if not user.id:
            return []
        tid = user.tenant_id
        if tid is not None and tid != 1:
            return [(getattr(self.model, "tenant_id") == tid) | (getattr(self.model, "tenant_id") == 1)]
        return []

    def _tenant_dml_where(self, sql):
        """为 DML 语句注入 tenant_id 条件（不读平台数据）"""
        if hasattr(self.model, "tenant_id"):
            user = self.auth.user
            if user.id and not user.is_superuser:
                tid = user.tenant_id
                if tid is not None:
                    return sql.where(getattr(self.model, "tenant_id") == tid)
        return sql

    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())

        if hasattr(self.model, "tenant_id") and not getattr(self.model, "__platform_data_shared__", False):
            user = self.auth.user
            if user.id and not user.is_superuser:
                tid = user.tenant_id
                if tid is not None:
                    conditions.append(getattr(self.model, "tenant_id") == tid)

        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))
            else:
                conditions.append(attr == value)
        return conditions

    @staticmethod
    def _resolve_condition(attr: ColumnElement, value: tuple) -> list[ColumnElement]:
        """解析 (operator, value) 元组为 SQLAlchemy 条件列表"""
        seq, val = value

        handlers: dict[str, tuple] = {
            "None": (lambda: [attr.is_(None)], True),
            "not None": (lambda: [attr.isnot(None)], True),
        }
        # 需要额外校验的运算符
        if seq in handlers:
            fn, _always = handlers[seq]
            return fn()

        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: dict[str, Any] = {
            "!=": 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]:
        """解析排序参数

        参数:
        - order: 排序字段列表, 格式为 [{'id': 'asc'}, {'name': 'desc'}]

        返回:
        - 排序表达式列表
        """
        columns: list[ColumnElement] = []
        for item in order:
            for field, direction in item.items():
                column = getattr(self.model, field)
                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]:
        """构建预加载选项

        参数:
        - preload: 预加载关系，支持关系名字符串或 SQLAlchemy loader option

        返回:
        - 预加载选项列表
        """
        options: list[Any] = []
        model_loader_options = getattr(self.model, "__loader_options__", [])

        all_preloads: set[str | Any] = set(model_loader_options)
        if preload:
            for opt in preload:
                if isinstance(opt, str):
                    all_preloads.add(opt)
        elif preload == []:
            all_preloads = set()

        for opt in all_preloads:
            if isinstance(opt, str):
                if hasattr(self.model, opt):
                    options.append(selectinload(getattr(self.model, opt)))
            else:
                options.append(opt)

        return options