hd-crane-gs/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, text
from sqlalchemy import inspect as sa_inspect

from app.core.base_model import MappedBase
from app.core.exceptions import CustomException
from app.core.permission import Permission
from app.api.v1.module_system.auth.schema import AuthSchema

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
    
    async def get(self, preload: Optional[List[Union[str, Any]]] = None, **kwargs) -> Optional[ModelType]:
        """
        根据条件获取单个对象
        
        参数:
        - preload (Optional[List[Union[str, Any]]]): 预加载关系，支持关系名字符串或SQLAlchemy loader option
        - **kwargs: 查询条件
            
        返回:
        - Optional[ModelType]: 对象实例
            
        异常:
        - CustomException: 查询失败时抛出异常
        """
        try:
            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.auth.db.execute(sql)
            obj = result.scalars().first()
            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, preload: Optional[List[Union[str, Any]]] = None) -> Sequence[ModelType]:
        """
        根据条件获取对象列表
        
        参数:
        - search (Optional[Dict]): 查询条件,格式为 {'id': value, 'name': value}
        - order_by (Optional[List[Dict[str, str]]]): 排序字段,格式为 [{'id': 'asc'}, {'name': 'desc'}]
        - preload (Optional[List[Union[str, Any]]]): 预加载关系，支持关系名字符串或SQLAlchemy loader option
            
        返回:
        - 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))
            # 应用可配置的预加载选项
            for opt in self.__loader_options(preload):
                sql = sql.options(opt)
            sql = await self.__filter_permissions(sql)
            result: Result = await self.auth.db.execute(sql)
            return result.scalars().all()
        except Exception as e:
            raise CustomException(msg=f"列表查询失败: {str(e)}")

    async def list_sql(self, sql:str,params: Optional[Dict[str, Any]] = None) -> List[Dict]:
        """
        根据条件获取对象列表

        参数:
        - search (Optional[Dict]): 查询条件,格式为 {'id': value, 'name': value}
        - order_by (Optional[List[Dict[str, str]]]): 排序字段,格式为 [{'id': 'asc'}, {'name': 'desc'}]
        - preload (Optional[List[Union[str, Any]]]): 预加载关系，支持关系名字符串或SQLAlchemy loader option

        返回:
        - Sequence[ModelType]: 对象列表

        异常:
        - CustomException: 查询失败时抛出异常
        """
        try:
            business_params = params.copy() if params else {}

            result = await self.execute_raw_sql(
                sql=sql,
                params=business_params,
                fetch_one=False,
                scalar=False
            )
            return result
        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', preload: Optional[List[Union[str, Any]]] = None) -> Sequence[ModelType]:
        """
        获取树形结构数据列表
        
        参数:
        - search (Optional[Dict]): 查询条件
        - order_by (Optional[List[Dict[str, str]]]): 排序字段
        - children_attr (str): 子节点属性名
        - preload (Optional[List[Union[str, Any]]]): 额外预加载关系，若为None则默认包含children_attr
            
        返回:
        - 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))
            
            # 处理预加载选项
            final_preload = preload
            # 如果没有提供preload且children_attr存在，则添加到预加载选项中
            if preload is None and children_attr and hasattr(self.model, children_attr):
                # 获取模型默认预加载选项
                model_defaults = getattr(self.model, "__loader_options__", [])
                # 将children_attr添加到默认预加载选项中
                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.auth.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], preload: Optional[List[Union[str, Any]]] = None) -> Dict:
        """
        获取分页数据
        
        参数:
        - offset (int): 偏移量
        - limit (int): 每页数量
        - order_by (List[Dict[str, str]]): 排序字段
        - search (Dict): 查询条件
        - out_schema (Type[OutSchemaType]): 输出数据模型
        - preload (Optional[List[Union[str, Any]]]): 预加载关系
            
        返回:
        - Dict: 分页数据
            
        异常:
        - 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))
            # 应用预加载选项
            for opt in self.__loader_options(preload):
                sql = sql.options(opt)
            sql = await self.__filter_permissions(sql)

            # 优化count查询：使用主键计数而非全表扫描
            mapper = sa_inspect(self.model)
            pk_cols = list(getattr(mapper, "primary_key", []))
            if pk_cols:
                # 使用主键的第一列进行计数（主键必定非NULL，性能更好）
                count_sql = select(func.count(pk_cols[0])).select_from(self.model)
            else:
                # 降级方案：使用count(*)
                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.auth.db.execute(count_sql)
            total = total_result.scalar() or 0

            result: Result = await self.auth.db.execute(sql.offset(offset).limit(limit))
            objs = result.scalars().all()

            return {
                "page_no": offset // limit + 1 if limit else 1,
                "page_size": limit if limit else 10,
                "total": total,
                "has_next": offset + limit < total,
                "items": [out_schema.model_validate(obj).model_dump() for obj in objs]
            }
        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)
            
            # 设置字段值（只检查一次current_user）
            if self.auth.user:
                if hasattr(obj, "created_id"):
                    setattr(obj, "created_id", self.auth.user.id)
                if hasattr(obj, "updated_id"):
                    setattr(obj, "updated_id", self.auth.user.id)
            
            self.auth.db.add(obj)
            await self.auth.db.flush()
            await self.auth.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="更新对象不存在")
            
            # 设置字段值（只检查一次current_user）
            if self.auth.user:
                if hasattr(obj, "updated_id"):
                    setattr(obj, "updated_id", self.auth.user.id)
            
            for key, value in obj_dict.items():
                if hasattr(obj, key):
                    setattr(obj, key, value)
                    
            await self.auth.db.flush()
            await self.auth.db.refresh(obj)
            
            # 权限二次确认：flush后再次验证对象仍在权限范围内
            # 防止并发修改导致的权限逃逸（如其他事务修改了created_id）
            verify_obj = await self.get(id=id)
            if not verify_obj:
                # 对象已被删除或权限已失效
                raise CustomException(msg="更新失败，对象不存在或无权限访问")
            
            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:
            # 先查询确认权限,避免删除无权限的数据
            objs = await self.list(search={"id": ("in", ids)})
            accessible_ids = [obj.id for obj in objs]
            
            # 检查是否所有ID都有权限访问
            inaccessible_count = len(ids) - len(accessible_ids)
            if inaccessible_count > 0:
                raise CustomException(msg=f"无权限删除{inaccessible_count}条数据")
            
            if not accessible_ids:
                return  # 没有可删除的数据
            
            mapper = sa_inspect(self.model)
            pk_cols = list(getattr(mapper, "primary_key", []))
            if not pk_cols:
                raise CustomException(msg="模型缺少主键，无法删除")
            if len(pk_cols) > 1:
                raise CustomException(msg="暂不支持复合主键的批量删除")
            
            # 只删除有权限的数据
            sql = delete(self.model).where(pk_cols[0].in_(accessible_ids))
            await self.auth.db.execute(sql)
            await self.auth.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.auth.db.execute(sql)
            await self.auth.db.flush()
        except Exception as e:
            raise CustomException(msg=f"清空失败: {str(e)}")

    async def execute_raw_sql(
            self,
            sql: str,
            params: Optional[Dict[str, Any]] = None,
            fetch_one: bool = False,
            scalar: bool = False
    ) -> Optional[Union[Dict, List[Dict], Any, Sequence[Any]]]:
        try:
            # ---------------------- 1. 严格校验输入：避免空SQL、非字符串等错误 ----------------------
            # 校验SQL是否为非空字符串
            if not isinstance(sql, str) or len(sql.strip()) == 0:
                raise CustomException(msg="传入的原始SQL不能为空且必须为字符串类型")

            # 初始化最终SQL（仅对字符串操作，规避TextClause错误）
            final_sql = sql.strip()
            # 初始化最终参数：兜底空字典，避免None报错
            final_params = params.copy() if (params and isinstance(params, Dict)) else {}

            # ---------------------- 3. 核心：执行SQL（仅包装一次TextClause，不进行后续字符串操作） ----------------------
            # 包装为TextClause（SQLAlchemy执行原始SQL必需，仅执行此步骤，无额外操作）
            raw_sql_clause = text(final_sql)

            # 执行SQL：传入上层已合并（业务+权限）的参数，确保所有占位符绑定完成
            db_result: Result = await self.auth.db.execute(raw_sql_clause, final_params)

            # ---------------------- 4. 处理结果：格式统一，兼容上层list_sql的返回需求 ----------------------
            # 处理SELECT查询（返回字典/字典列表，兼容Sequence[ModelType]）
            if final_sql.upper().startswith("SELECT"):
                # 标量结果（单个值，如COUNT(*)）
                if scalar:
                    return db_result.scalar() if fetch_one else db_result.scalars().all()

                # 完整行结果：转换为字典列表，避免字段名丢失，兜底空列表
                row_mappings = db_result.mappings().all()
                result_list = [dict(row) for row in row_mappings] if row_mappings else []

                # 单条/多条结果返回，确保返回序列类型
                if fetch_one:
                    return result_list[0] if len(result_list) > 0 else None
                else:
                    return result_list

            # 处理DML操作（INSERT/UPDATE/DELETE）：刷新会话，返回None
            await self.auth.db.flush()
            return None

        # ---------------------- 5. 异常捕获：补充上下文，方便排查上层拼接后的问题 ----------------------
        except CustomException as ce:
            raise ce
        except Exception as e:
            error_msg = (
                f"执行原始SQL失败：{str(e)}"
                f"\n  执行的SQL：{final_sql[:500]}..."
                f"\n  传入的参数：{final_params}"
            )
            raise CustomException(msg=error_msg)

    async def set(self, ids: List[int], **kwargs) -> None:
        """
        批量更新对象
        
        参数:
        - ids (List[int]): 对象ID列表
        - **kwargs: 更新的属性及值
            
        异常:
        - CustomException: 更新失败时抛出异常
        """
        try:
            # 先查询确认权限,避免更新无权限的数据
            objs = await self.list(search={"id": ("in", ids)})
            accessible_ids = [obj.id for obj in objs]
            
            # 检查是否所有ID都有权限访问
            inaccessible_count = len(ids) - len(accessible_ids)
            if inaccessible_count > 0:
                raise CustomException(msg=f"无权限更新{inaccessible_count}条数据")
            
            if not accessible_ids:
                return  # 没有可更新的数据
            
            mapper = sa_inspect(self.model)
            pk_cols = list(getattr(mapper, "primary_key", []))
            if not pk_cols:
                raise CustomException(msg="模型缺少主键，无法更新")
            if len(pk_cols) > 1:
                raise CustomException(msg="暂不支持复合主键的批量更新")
            
            # 只更新有权限的数据
            sql = update(self.model).where(pk_cols[0].in_(accessible_ids)).values(**kwargs)
            await self.auth.db.execute(sql)
            await self.auth.db.flush()
        except CustomException:
            raise
        except Exception as e:
            raise CustomException(msg=f"批量更新失败: {str(e)}")

    async def __filter_permissions(self, sql: Select) -> Select:
        """
        过滤数据权限（仅用于Select）。
        """
        filter = Permission(
            model=self.model,
            auth=self.auth
        )
        return await filter.filter_query(sql)


    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 == ">" and val:
                    conditions.append(attr > val)
                elif seq == ">=" and val:
                    conditions.append(attr >= val)
                elif seq == "<" and val:
                    conditions.append(attr < val)
                elif seq == "<=" and val:
                    conditions.append(attr <= val)
                elif seq == "==" and val:
                    conditions.append(attr == val)
            else:
                conditions.append(attr == value)
        return conditions

    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

    def __loader_options(self, preload: Optional[List[Union[str, Any]]] = None) -> List[Any]:
        """
        构建预加载选项
        
        参数:
        - preload (Optional[List[Union[str, Any]]]): 预加载关系，支持关系名字符串或SQLAlchemy loader option
            
        返回:
        - List[Any]: 预加载选项列表
        """
        options = []
        # 获取模型定义的默认加载选项
        model_loader_options = getattr(self.model, '__loader_options__', [])
        
        # 合并所有需要预加载的选项
        all_preloads = 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):
                # 使用selectinload来避免在异步环境中的MissingGreenlet错误
                if hasattr(self.model, opt):
                    options.append(selectinload(getattr(self.model, opt)))
            else:
                # 直接使用非字符串的加载选项
                options.append(opt)
                
        return options