FastapiAdmin/backend/app/core/base_crud.py

# -*- coding: utf-8 -*-

from pydantic import BaseModel
from typing import TypeVar, Sequence, Generic, Dict, Any, List, Optional, Type, Union
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.orm import selectinload
from sqlalchemy.engine import Result
from sqlalchemy import asc, func, select, delete, Select, desc, update, or_, and_

from app.core.base_model import MappedBase
from app.api.v1.module_system.auth.schema import AuthSchema
from app.api.v1.module_system.dept.model import DeptModel
from app.api.v1.module_system.user.model import UserModel
from app.utils.common_util import get_child_id_map, get_child_recursion
from app.core.exceptions import CustomException
from app.common.request import PageResultSchema
from app.core.serialize import Serialize

ModelType = TypeVar("ModelType", bound=MappedBase)
CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)
OutSchemaType = TypeVar("OutSchemaType", bound=BaseModel)


class CRUDBase(Generic[ModelType, CreateSchemaType, UpdateSchemaType]):
    """基础数据层"""

    def __init__(self, model: Type[ModelType], auth: AuthSchema) -> None:
        """
        初始化CRUDBase类
        
        参数:
        - model (Type[ModelType]): 数据模型类。
        - auth (AuthSchema): 认证信息。

        返回:
        - None
        """
        self.model = model
        self.auth = auth
        self.db = auth.db
        self.current_user = auth.user
    
    async def get(self, **kwargs) -> Optional[ModelType]:
        """
        根据条件获取单个对象
        
        参数:
        - **kwargs: 查询条件
            
        返回:
        - Optional[ModelType]: 对象实例
            
        返回:
        - Optional[ModelType]: 对象实例
            
        异常:
        - CustomException: 查询失败时抛出异常
        """
        try:
            conditions = await self.__build_conditions(**kwargs)
            sql = select(self.model).where(*conditions)
            # 只有继承自CreatorMixin的模型才有creator关系
            if hasattr(self.model, "creator_id"):
                sql = sql.options(selectinload(self.model.creator))
            
            sql = await self.__filter_permissions(sql)

            result: Result = await self.db.execute(sql)
            obj = result.scalars().first()
            # if not obj:
            #     raise CustomException(msg="该信息不存在")

            return obj
        except Exception as e:
            raise CustomException(msg=f"获取查询失败: {str(e)}")

    async def list(self, search: Optional[Dict] = None, order_by: Optional[List[Dict[str, str]]] = None) -> Sequence[ModelType]:
        """
        根据条件获取对象列表和总数
        
        参数:
        - search (Optional[Dict]): 查询条件,格式为 {'id': value, 'name': value}
        - order_by (Optional[List[Dict[str, str]]]): 排序字段,格式为 [{'id': 'asc'}, {'name': 'desc'}]
            
        返回:
        - Sequence[ModelType]: 对象列表和总数
            
        异常:
        - CustomException: 查询失败时抛出异常
        """
        try:
            conditions = await self.__build_conditions(**search) if search else []
            order = order_by or [{'id': 'asc'}]
            sql = select(self.model).where(*conditions).order_by(*self.__order_by(order))
            # 只有继承自CreatorMixin的模型才有creator关系
            if hasattr(self.model, "creator_id"):
                sql = sql.options(selectinload(self.model.creator))
            sql = await self.__filter_permissions(sql)
            result: Result = await self.db.execute(sql)
            return result.scalars().all()
        except Exception as e:
            raise CustomException(msg=f"列表查询失败: {str(e)}")

    async def tree_list(self, search: Optional[Dict] = None, order_by: Optional[List[Dict[str, str]]] = None, children_attr: str = 'children') -> Sequence[ModelType]:
        """
        获取树形结构数据列表
        
        参数:
        - search (Optional[Dict]): 查询条件
        - order_by (Optional[List[Dict[str, str]]]): 排序字段
        - children_attr (str): 子节点属性名
            
        返回:
        - Sequence[ModelType]: 树形结构数据列表
            
        异常:
        - CustomException: 查询失败时抛出异常
        """
        try:
            from sqlalchemy.orm import selectinload
            
            conditions = await self.__build_conditions(**search) if search else []
            order = order_by or [{'id': 'asc'}]
            sql = select(self.model).where(*conditions).order_by(*self.__order_by(order))
            
            # 如果模型有children属性，则预加载该关系
            if hasattr(self.model, children_attr):
                sql = sql.options(selectinload(getattr(self.model, children_attr)))
            
            # 只有继承自CreatorMixin的模型才有creator关系
            if hasattr(self.model, "creator_id"):
                sql = sql.options(selectinload(self.model.creator))
                
            sql = await self.__filter_permissions(sql)
            result: Result = await self.db.execute(sql)
            return result.scalars().all()
        except Exception as e:
            raise CustomException(msg=f"树形列表查询失败: {str(e)}")
    
    async def page(self, offset: int, limit: int, order_by: List[Dict[str, str]], search: Dict, out_schema: Type[OutSchemaType]) -> Dict:
        """
        获取分页数据
        
        参数:
        - offset (int): 偏移量
        - limit (int): 每页数量
        - order_by (List[Dict[str, str]]): 排序字段
        - search (Dict): 查询条件
        - out_schema (Type[OutSchemaType]): 输出数据模型
            
        返回:
        - Dict: 分页数据
            
        异常:
        - CustomException: 查询失败时抛出异常
        """
        try:
            from sqlalchemy.orm import selectinload

            conditions = await self.__build_conditions(**search) if search else []
            order = order_by or [{'id': 'asc'}]
            sql = select(self.model).where(*conditions).order_by(*self.__order_by(order))
            # 只有继承自CreatorMixin的模型才有creator关系
            if hasattr(self.model, "creator_id"):
                sql = sql.options(selectinload(self.model.creator))
            sql = await self.__filter_permissions(sql)

            # 获取总数
            count_sql = select(func.count()).select_from(self.model)
            # 应用相同的过滤条件到计数查询
            if conditions:
                count_sql = count_sql.where(*conditions)
            count_sql = await self.__filter_permissions(count_sql)
            
            total_result = await self.db.execute(count_sql)
            total = total_result.scalar()
            
            if total is None:
                total = 0

            result: Result = await self.db.execute(sql.offset(offset).limit(limit))

            objs = result.scalars().all()

            data=PageResultSchema(
                items=[out_schema.model_validate(obj).model_dump()  for obj in objs],
                total=total,
                page_no=offset // limit + 1 if limit else 1,
                page_size=limit,
                has_next=offset + limit < total,
            ).model_dump()

            return data
        except Exception as e:
            raise CustomException(msg=f"分页查询失败: {str(e)}")
    
    async def create(self, data: Union[CreateSchemaType, Dict]) -> ModelType:
        """
        创建新对象
        
        参数:
        - data (Union[CreateSchemaType, Dict]): 对象属性
            
        返回:
        - ModelType: 新创建的对象实例
            
        异常:
        - CustomException: 创建失败时抛出异常
        """
        try:
            obj_dict = data if isinstance(data, dict) else data.model_dump()
            obj = self.model(**obj_dict)
            
            # 只有继承自CreatorMixin的模型才有creator关系
            if hasattr(self.model, "creator_id") and self.current_user:
                # 设置创建人ID
                obj.creator_id = self.current_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"创建失败: {str(e)}")

    async def update(self, id: int, data: Union[UpdateSchemaType, Dict]) -> ModelType:
        """
        更新对象
        
        参数:
        - id (int): 对象ID
        - data (Union[UpdateSchemaType, Dict]): 更新的属性及值
            
        返回:
        - ModelType: 更新后的对象实例
            
        异常:
        - CustomException: 更新失败时抛出异常
        """
        try:
            obj_dict = data if isinstance(data, dict) else data.model_dump(exclude_unset=True, exclude={"id"})
            obj = await self.get(id=id)
            if not obj:
                raise CustomException(msg="更新对象不存在")
            
            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 Exception as e:
            raise CustomException(msg=f"更新失败: {str(e)}")

    async def delete(self, ids: List[int]) -> None:
        """
        删除对象
        
        参数:
        - ids (List[int]): 对象ID列表
            
        异常:
        - CustomException: 删除失败时抛出异常
        """
        try:
            sql = delete(self.model).where(self.model.id.in_(ids))
            sql = await self.__filter_permissions(sql)
            await self.db.execute(sql)
            await self.db.flush()
        except Exception as e:
            raise CustomException(msg=f"删除失败: {str(e)}")

    async def clear(self) -> None:
        """
        清空对象表
        
        异常:
        - CustomException: 清空失败时抛出异常
        """
        try:
            sql = delete(self.model)
            await self.db.execute(sql)
            await self.db.flush()
        except Exception as e:
            raise CustomException(msg=f"清空失败: {str(e)}")

    async def set(self, ids: List[int], **kwargs) -> None:
        """
        批量更新对象
        
        参数:
        - ids (List[int]): 对象ID列表
        - **kwargs: 更新的属性及值
            
        异常:
        - CustomException: 更新失败时抛出异常
        """
        try:
            sql = update(self.model).where(self.model.id.in_(ids)).values(**kwargs)
            await self.db.execute(sql)
            await self.db.flush()
        except Exception as e:
            raise CustomException(msg=f"批量更新失败: {str(e)}")

    async def __filter_permissions(self, sql: Select) -> Select:
        """
        过滤数据权限
        
        参数:
        - sql (Select): SQL查询对象
            
        返回:
        - Select: 过滤后的数据查询对象
            
        异常:
        - CustomException: 权限过滤失败时抛出异常
        """
        # 如果不需要检查数据权限,则直接返回
        if not self.current_user or not self.auth.check_data_scope:
            return sql

        # 1. 如果模型没有创建人creator字段,则不需要权限判断
        if not hasattr(self.model, "creator_id"):
            return sql
        
        sql = sql.options(selectinload(self.model.creator))
        
        # 2. 超级管理员可以查看所有数据
        if self.current_user.is_superuser:
            return sql
            
        # 3. 如果用户没有部门或角色,则只能查看自己的数据
        if not self.current_user.dept_id or not self.current_user.roles:
            return sql.where(self.model.creator_id == self.current_user.id)
        
        # 4. 获取用户所有角色的权限范围
        data_scopes = set()
        dept_ids = set()
        
        # data_scope 数据权限范围说明:
        # 1: 仅本人数据权限
        # 2: 本部门数据权限  
        # 3: 本部门及以下数据权限
        # 4: 全部数据权限
        # 5: 自定义数据权限
        
        # 获取当前用户所绑定角色的数据权限范围
        for role in self.current_user.roles:
            # 检查role是否有depts属性
            if hasattr(role, 'depts'):
                for dept in role.depts:
                    dept_ids.add(dept.id)

            data_scopes.add(role.data_scope)
        
        if 4 in data_scopes:
            # 4、全部数据权限
            return sql

        if 1 in data_scopes:
            # 1、仅本人数据
            return sql.where(self.model.creator_id == self.current_user.id)

        if 2 in data_scopes:
            # 2、本部门数据
            dept_ids.add(self.current_user.dept_id)
            
        if 3 in data_scopes:
            # 3、本部门及以下数据
            # 直接查询部门表，避免递归调用CRUD
            dept_sql = select(DeptModel)
            dept_result = await self.db.execute(dept_sql)
            dept_objs = dept_result.scalars().all()
            id_map = get_child_id_map(dept_objs)
            dept_child_ids = get_child_recursion(id=self.current_user.dept_id, id_map=id_map)
            for child_id in dept_child_ids:
                dept_ids.add(child_id)

        # 5、自定义权限
        # 检查UserModel是否有dept_id属性
        if hasattr(UserModel, 'dept_id'):
            return sql.where(self.model.creator.has(UserModel.dept_id.in_(list(dept_ids))))
        else:
            # 如果没有dept_id属性，回退到只显示自己的数据
            return sql.where(self.model.creator_id == self.current_user.id)

    def __order_by(self, order_by: List[Dict[str, str]]) -> List[ColumnElement]:
        """
        获取排序字段
        
        参数:
        - order_by (List[Dict[str, str]]): 排序字段列表,格式为 [{'id': 'asc'}, {'name': 'desc'}]
            
        返回:
        - List[ColumnElement]: 排序字段列表
            
        异常:
        - CustomException: 排序字段不存在时抛出异常
        """
        columns = []
        for order in order_by:
            for field, direction in order.items():
                column = getattr(self.model, field)
                columns.append(desc(column) if direction.lower() == 'desc' else asc(column))
        return columns

    async def __build_conditions(self, **kwargs) -> List[ColumnElement]:
        """
        构建查询条件
        
        参数:
        - **kwargs: 查询参数
            
        返回:
        - List[ColumnElement]: SQL条件表达式列表
            
        异常:
        - CustomException: 查询参数不存在时抛出异常
        """
        conditions = []
        for key, value in kwargs.items():
            if value is None or value == "":
                continue

            attr = getattr(self.model, key)
            if isinstance(value, tuple):
                seq, val = value
                if seq == "None":
                    conditions.append(attr.is_(None))
                elif seq == "not None":
                    conditions.append(attr.isnot(None))
                elif seq == "date" and val:
                    conditions.append(func.date_format(attr, "%Y-%m-%d") == val)
                elif seq == "month" and val:
                    conditions.append(func.date_format(attr, "%Y-%m") == val)
                elif seq == "like" and val:
                    conditions.append(attr.like(f"%{val}%"))
                elif seq == "in" and val:
                    conditions.append(attr.in_(val))
                elif seq == "between" and isinstance(val, (list, tuple)) and len(val) == 2:
                    conditions.append(attr.between(val[0], val[1]))
                elif seq == "!=" and val:
                    conditions.append(attr != val)
                elif seq in [">", ">=", "<=", "=="] and val:
                    conditions.append(getattr(attr, seq.replace("==", "__eq__"))(val))
            else:
                conditions.append(attr == value)
        return conditions