|
|
import functools
|
|
|
import inspect
|
|
|
from contextlib import asynccontextmanager
|
|
|
from typing import Any, ParamSpec, TypeVar, Callable
|
|
|
|
|
|
from sqlalchemy import Executable, Result, and_, Select, Delete, Update
|
|
|
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
|
|
from sqlalchemy.sql.selectable import Subquery
|
|
|
|
|
|
from common.constant import Constant
|
|
|
from common.global_enums import IsDelete
|
|
|
from config import settings
|
|
|
|
|
|
P = ParamSpec('P')
|
|
|
T = TypeVar('T')
|
|
|
|
|
|
engine = create_async_engine(settings.database_url, echo=False, # 打印SQL日志（生产环境建议关闭）
|
|
|
pool_size=10, # 连接池大小
|
|
|
max_overflow=20, # 最大溢出连接数
|
|
|
pool_recycle=3600, # 连接回收时间（秒），解决MySQL超时断开问题【4†source】【5†source】
|
|
|
)
|
|
|
|
|
|
|
|
|
# 创建异步会话工厂
|
|
|
|
|
|
|
|
|
class EnhanceAsyncSession(AsyncSession):
|
|
|
|
|
|
def _add_logical_delete_condition(self, statement: Select) -> Select:
|
|
|
"""
|
|
|
为 Select 语句添加逻辑删除条件
|
|
|
支持递归处理 fastapi-pagination 生成的子查询包装
|
|
|
"""
|
|
|
# 特殊处理：如果当前查询只有一个子查询，直接修改其内部元素
|
|
|
# 这样可以避免产生重复的别名
|
|
|
if len(statement.froms) == 1 and isinstance(statement.froms[0], Subquery):
|
|
|
subquery = statement.froms[0]
|
|
|
# 递归处理子查询内部
|
|
|
processed_inner = self._add_logical_delete_condition(subquery.element)
|
|
|
# 如果内部有修改，直接替换 subquery 的 element
|
|
|
if processed_inner is not subquery.element:
|
|
|
subquery.element = processed_inner
|
|
|
return statement
|
|
|
|
|
|
# 处理当前层的表
|
|
|
delete_condition = None
|
|
|
delete_field = Constant.LOGICAL_DELETE_FIELD
|
|
|
|
|
|
for from_obj in statement.froms:
|
|
|
# 跳过子查询（因为会通过上面的逻辑递归处理）
|
|
|
if isinstance(from_obj, Subquery):
|
|
|
continue
|
|
|
|
|
|
# 只处理 Table 对象
|
|
|
if hasattr(from_obj, 'columns') and delete_field in from_obj.columns:
|
|
|
# 使用表对象上的列
|
|
|
condition = from_obj.columns[delete_field] == IsDelete.NO_DELETE
|
|
|
|
|
|
if delete_condition is None:
|
|
|
delete_condition = condition
|
|
|
else:
|
|
|
delete_condition = and_(delete_condition, condition)
|
|
|
|
|
|
# 如果有条件，则应用到当前 statement
|
|
|
if delete_condition is not None:
|
|
|
existing_condition = statement.whereclause
|
|
|
if existing_condition is not None:
|
|
|
new_condition = and_(existing_condition, delete_condition)
|
|
|
else:
|
|
|
new_condition = delete_condition
|
|
|
|
|
|
statement = statement.where(new_condition)
|
|
|
|
|
|
return statement
|
|
|
|
|
|
async def scalar(self, statement: Executable, params=None, *, execution_options=None, bind_arguments=None,
|
|
|
**kw: Any, ):
|
|
|
sig = inspect.signature(super().scalar)
|
|
|
if execution_options is None:
|
|
|
default_execution_options = sig.parameters['execution_options'].default
|
|
|
execution_options = default_execution_options
|
|
|
|
|
|
# 只对 Select 语句添加逻辑删除条件
|
|
|
if isinstance(statement, Select):
|
|
|
statement = self._add_logical_delete_condition(statement)
|
|
|
|
|
|
return await super().scalar(statement, params, execution_options=execution_options,
|
|
|
bind_arguments=bind_arguments, **kw)
|
|
|
|
|
|
async def execute(self, statement: Executable, params=None, *, execution_options=None, bind_arguments=None,
|
|
|
**kw: Any, ) -> Result[Any]:
|
|
|
sig = inspect.signature(super().execute)
|
|
|
if execution_options is None:
|
|
|
default_execution_options = sig.parameters['execution_options'].default
|
|
|
execution_options = default_execution_options
|
|
|
|
|
|
if isinstance(statement, Select):
|
|
|
statement = self._add_logical_delete_condition(statement)
|
|
|
|
|
|
if isinstance(statement, Delete):
|
|
|
skip_soft_delete = execution_options and execution_options.get("skip_soft_delete", False)
|
|
|
|
|
|
if not skip_soft_delete:
|
|
|
table = statement.table
|
|
|
|
|
|
if hasattr(table, 'columns') and Constant.LOGICAL_DELETE_FIELD in table.columns:
|
|
|
update_stmt = (Update(table).where(statement.whereclause).values(
|
|
|
**{Constant.LOGICAL_DELETE_FIELD: IsDelete.DELETE}))
|
|
|
|
|
|
if statement._returning:
|
|
|
update_stmt = update_stmt.returning(*statement._returning)
|
|
|
|
|
|
return await super().execute(update_stmt, params=params, execution_options=execution_options,
|
|
|
bind_arguments=bind_arguments, **kw)
|
|
|
|
|
|
result = await super().execute(statement, params=params, execution_options=execution_options,
|
|
|
bind_arguments=bind_arguments, **kw, )
|
|
|
return result
|
|
|
|
|
|
def delete(self, instance):
|
|
|
from sqlalchemy import inspect as sqlalchemy_inspect
|
|
|
|
|
|
if hasattr(instance, Constant.LOGICAL_DELETE_FIELD):
|
|
|
setattr(instance, Constant.LOGICAL_DELETE_FIELD, IsDelete.DELETE)
|
|
|
|
|
|
insp = sqlalchemy_inspect(instance)
|
|
|
if insp.detached or insp.transient:
|
|
|
self.add(instance)
|
|
|
else:
|
|
|
super().delete(instance)
|
|
|
|
|
|
|
|
|
AsyncSessionLocal = async_sessionmaker(bind=engine, class_=EnhanceAsyncSession, expire_on_commit=False, # 提交后不使对象过期
|
|
|
autoflush=False # 禁用自动刷新
|
|
|
)
|
|
|
|
|
|
|
|
|
# 获取数据库session的独立方法
|
|
|
@asynccontextmanager
|
|
|
async def get_db_session():
|
|
|
"""获取数据库session的上下文管理器"""
|
|
|
async with AsyncSessionLocal() as session:
|
|
|
async with session.begin():
|
|
|
yield session
|
|
|
|
|
|
|
|
|
def with_db_session(session_param_name: str = "session"):
|
|
|
"""
|
|
|
一个装饰器，用于为异步函数自动注入数据库会话。
|
|
|
|
|
|
Args:
|
|
|
session_param_name: 被装饰函数中，用于接收会话的参数名，默认为 'session'。
|
|
|
"""
|
|
|
|
|
|
def decorator(func: Callable[P, T]) -> Callable[P, T]:
|
|
|
# 确保只装饰异步函数
|
|
|
if not inspect.iscoroutinefunction(func):
|
|
|
raise TypeError("`with_db_session` can only be used on async functions.")
|
|
|
|
|
|
@functools.wraps(func)
|
|
|
async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
|
|
|
# 如果调用时已经手动传了 session，就直接用
|
|
|
if session_param_name in kwargs:
|
|
|
return await func(*args, **kwargs)
|
|
|
|
|
|
# 否则，创建一个新 session 并注入
|
|
|
async with get_db_session() as session:
|
|
|
kwargs[session_param_name] = session
|
|
|
return await func(*args, **kwargs)
|
|
|
|
|
|
return wrapper
|
|
|
|
|
|
return decorator
|
|
|
|
|
|
|
|
|
# 关闭引擎
|
|
|
async def close_engine():
|
|
|
await engine.dispose()
|