import asyncio import logging from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from dataclasses import dataclass, field from datetime import datetime, timedelta from functools import cache from typing import Any, TypeVar, cast from pydantic import SecretStr from sqlalchemy import URL, Engine, create_engine, text from sqlalchemy.exc import DBAPIError from sqlalchemy.ext.asyncio import ( AsyncAttrs, AsyncConnection, AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine, ) from sqlalchemy.orm import DeclarativeBase, sessionmaker from config.env import DataBaseConfig, DataBaseSettings, DataSourceSettings from exceptions.exception import ( DataSourceException, DataSourceInitializationException, DataSourceNotFoundException, DataSourceUnavailableException, ) from utils.log_util import logger from utils.time_util import TimezoneUtil _HEALTH_RETRY_COOLDOWN = timedelta(seconds=5) _SESSION_TIMEZONE_QUERY = { 'mysql': 'SELECT @@session.time_zone', 'postgresql': 'SHOW TIME ZONE', } _UTC_SESSION_TIMEZONES = {'+00:00', 'UTC', 'Etc/UTC'} _QUIET_SQL_LOG_TOKEN = 'ruoyi_internal_sql' _DatabaseEngine = TypeVar('_DatabaseEngine', Engine, AsyncEngine) class _QuietSqlLogFilter(logging.Filter): """ 按连接日志标记过滤内部SQL，保留警告和错误 """ def filter(self, record: logging.LogRecord) -> bool: """ 判断是否保留当前SQLAlchemy日志 :param record: 待输出的日志记录 :return: 是否保留日志 """ return ( record.levelno >= logging.WARNING or not isinstance(record.msg, str) or not record.msg.startswith(f'[{_QUIET_SQL_LOG_TOKEN}] ') ) _QUIET_SQL_LOG_FILTER = _QuietSqlLogFilter() def quiet_sql_engine(engine: _DatabaseEngine) -> _DatabaseEngine: """ 创建复用原连接池、只过滤自身常规SQL日志的引擎视图 :param engine: 注册中心管理的同步或异步引擎 :return: 带有内部SQL日志标记的引擎视图，由原引擎统一管理连接池 """ sync_engine = engine.sync_engine if isinstance(engine, AsyncEngine) else engine logger_name = 'sqlalchemy.engine.Engine' if sync_engine.logging_name: logger_name = f'{logger_name}.{sync_engine.logging_name}' # Logger过滤同时覆盖echo自动创建的Handler和应用日志转发，不修改共享日志等级。 logging.getLogger(logger_name).addFilter(_QUIET_SQL_LOG_FILTER) return cast('_DatabaseEngine', engine.execution_options(logging_token=_QUIET_SQL_LOG_TOKEN)) def _error_details(exc: BaseException) -> tuple[str, int | None]: """ 提取不含连接凭据和SQL参数的安全错误摘要 :param exc: 原始异常或数据源异常 :return: 异常类型和数字错误码 """ if isinstance(exc, DataSourceException): return exc.error_type or type(exc).__name__, exc.error_code original = exc.orig if isinstance(exc, DBAPIError) else exc error_code = original.args[0] if original.args and isinstance(original.args[0], int) else None return type(original).__name__, error_code def _error_log_suffix(error_type: str, error_code: int | None) -> str: """ 构建数据源错误日志摘要 :param error_type: 异常类型 :param error_code: 数字错误码 :return: 不含敏感信息的日志后缀 """ code_text = f'，错误码：{error_code}' if error_code is not None else '' return f'，错误类型：{error_type}{code_text}' @dataclass(frozen=True, slots=True) class DatabaseDriverAdapter: """ 数据库驱动适配器 """ db_type: str async_driver: str sync_driver: str async_connect_timeout_key: str sync_connect_timeout_key: str def build_url(self, config: DataSourceSettings, *, sync: bool) -> URL: """ 根据数据源配置构建SQLAlchemy数据库连接URL :param config: 数据源配置 :param sync: 是否构建同步数据库连接URL :return: SQLAlchemy数据库连接URL """ return URL.create( drivername=self.sync_driver if sync else self.async_driver, username=config.db_username, password=_secret_value(config.db_password), host=config.db_host, port=int(config.db_port), database=config.db_database, ) def build_connect_args(self, config: DataSourceSettings, *, sync: bool) -> dict[str, Any]: """ 构建数据库驱动连接参数 :param config: 数据源配置 :param sync: 是否构建同步数据库连接参数 :return: 数据库驱动连接参数 """ timeout_key = self.sync_connect_timeout_key if sync else self.async_connect_timeout_key connect_args: dict[str, Any] = {timeout_key: config.db_connect_timeout} if self.db_type == 'postgresql': if sync: connect_args['options'] = '-c timezone=UTC' else: connect_args['server_settings'] = {'timezone': 'UTC'} elif self.db_type == 'mysql': connect_args['init_command'] = "SET time_zone = '+00:00'" return connect_args _DATABASE_DRIVER_ADAPTERS = { adapter.db_type: adapter for adapter in ( DatabaseDriverAdapter( db_type='mysql', async_driver='mysql+asyncmy', sync_driver='mysql+pymysql', async_connect_timeout_key='connect_timeout', sync_connect_timeout_key='connect_timeout', ), DatabaseDriverAdapter( db_type='postgresql', async_driver='postgresql+asyncpg', sync_driver='postgresql+psycopg2', async_connect_timeout_key='timeout', sync_connect_timeout_key='connect_timeout', ), ) } def _secret_value(value: SecretStr | str) -> str: """ 获取密码配置的原始值 :param value: 密码配置 :return: 密码原始值 """ return value.get_secret_value() if isinstance(value, SecretStr) else value def _database_source(config: DataBaseSettings | DataSourceSettings) -> DataSourceSettings: """ 获取指定配置对应的数据源配置 :param config: 数据库集合配置或单个数据源配置 :return: 单个数据源配置 """ return config.get_source() if isinstance(config, DataBaseSettings) else config def _driver_adapter(config: DataSourceSettings) -> DatabaseDriverAdapter: """ 获取数据库驱动适配器 :param config: 数据源配置 :return: 数据库驱动适配器 """ adapter = _DATABASE_DRIVER_ADAPTERS.get(config.db_type) if adapter is None: raise ValueError(f'不支持的数据库类型：{config.db_type!r}') return adapter def _build_url(config: DataSourceSettings, *, sync: bool) -> URL: """ 使用对应的数据库驱动适配器构建连接URL :param config: 数据源配置 :param sync: 是否构建同步数据库连接URL :return: SQLAlchemy数据库连接URL """ return _driver_adapter(config).build_url(config, sync=sync) def build_async_sqlalchemy_database_url(config: DataBaseSettings | DataSourceSettings | None = None) -> URL: """ 构建异步SQLAlchemy数据库连接URL :param config: 数据库集合配置或单个数据源配置 :return: 异步SQLAlchemy数据库连接URL """ return _build_url(_database_source(config or DataBaseConfig), sync=False) def build_sync_sqlalchemy_database_url(config: DataBaseSettings | DataSourceSettings | None = None) -> URL: """ 构建同步SQLAlchemy数据库连接URL :param config: 数据库集合配置或单个数据源配置 :return: 同步SQLAlchemy数据库连接URL """ return _build_url(_database_source(config or DataBaseConfig), sync=True) def _engine_options(config: DataSourceSettings, echo: bool | None = None) -> dict[str, Any]: """ 构建数据库引擎连接池参数 :param config: 数据源配置 :param echo: 是否输出SQLAlchemy SQL日志 :return: 数据库引擎连接池参数 """ return { 'echo': config.db_echo if echo is None else echo, 'max_overflow': config.db_max_overflow, 'pool_size': config.db_pool_size, 'pool_recycle': config.db_pool_recycle, 'pool_timeout': config.db_pool_timeout, 'pool_pre_ping': True, 'pool_use_lifo': True, } def create_async_db_engine( echo: bool | None = None, config: DataBaseSettings | DataSourceSettings | None = None ) -> AsyncEngine: """ 创建异步SQLAlchemy Engine :param echo: 是否输出SQLAlchemy SQL日志 :param config: 数据库集合配置或单个数据源配置 :return: 异步SQLAlchemy Engine """ source = _database_source(config or DataBaseConfig) adapter = _driver_adapter(source) return create_async_engine( adapter.build_url(source, sync=False), connect_args=adapter.build_connect_args(source, sync=False), **_engine_options(source, echo), ) def create_sync_db_engine( echo: bool | None = None, config: DataBaseSettings | DataSourceSettings | None = None ) -> Engine: """ 创建同步SQLAlchemy Engine :param echo: 是否输出SQLAlchemy SQL日志 :param config: 数据库集合配置或单个数据源配置 :return: 同步SQLAlchemy Engine """ source = _database_source(config or DataBaseConfig) adapter = _driver_adapter(source) return create_engine( adapter.build_url(source, sync=True), connect_args=adapter.build_connect_args(source, sync=True), **_engine_options(source, echo), ) def create_async_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]: """ 创建异步Session工厂 :param engine: 异步SQLAlchemy Engine :return: 异步Session工厂 """ return async_sessionmaker(bind=engine, autocommit=False, autoflush=False, expire_on_commit=False) def create_sync_session_factory(engine: Engine) -> sessionmaker: """ 创建同步Session工厂 :param engine: 同步SQLAlchemy Engine :return: 同步Session工厂 """ return sessionmaker(bind=engine, autocommit=False, autoflush=False) @dataclass(slots=True) class DataSourceRuntime: """ 数据源运行时状态 """ name: str config: DataSourceSettings async_engine: AsyncEngine | None = None async_session_factory: async_sessionmaker[AsyncSession] | None = None sync_engine: Engine | None = None available: bool = False last_health_check_at: datetime | None = None next_retry_at: datetime | None = None session_timezone: str | None = None health_lock: asyncio.Lock = field(default_factory=asyncio.Lock) class _DataSourceRegistry: """ 数据源注册中心 """ def __init__(self, settings: DataBaseSettings | None = None) -> None: self.settings = settings or DataBaseConfig self._runtimes: dict[str, DataSourceRuntime] = {} self._initialized = False self._log_enabled = True self._configs = dict(self.settings.db_sources) def _resolve_name(self, name: str | None = None) -> str: """ 解析并校验数据源名称 :param name: 数据源名称 :return: 已配置的数据源名称 """ source_name = name or self.settings.db_default_source if source_name not in self._configs: raise DataSourceNotFoundException(source_name) return source_name def _runtime(self, name: str | None = None) -> DataSourceRuntime: """ 获取或创建数据源运行时状态 :param name: 数据源名称 :return: 数据源运行时状态 """ source_name = self._resolve_name(name) runtime = self._runtimes.get(source_name) if runtime is None: runtime = DataSourceRuntime(name=source_name, config=self._configs[source_name]) self._runtimes[source_name] = runtime return runtime @staticmethod def _mark_unavailable(runtime: DataSourceRuntime) -> None: """ 标记数据源不可用并设置下次重试时间 :param runtime: 数据源运行时状态 :return: None """ now = TimezoneUtil.utc_now() runtime.available = False runtime.session_timezone = None runtime.last_health_check_at = now runtime.next_retry_at = now + _HEALTH_RETRY_COOLDOWN @staticmethod def _data_source_error( exception_type: type[DataSourceException], runtime: DataSourceRuntime, exc: BaseException, ) -> DataSourceException: """将底层异常转换为不泄露连接信息的数据源异常。""" error_type, error_code = _error_details(exc) return exception_type(runtime.name, error_type=error_type, error_code=error_code) @classmethod def _ensure_async_resources(cls, runtime: DataSourceRuntime) -> None: """ 确保数据源异步引擎和Session工厂已创建 :param runtime: 数据源运行时状态 :return: None """ if runtime.async_engine is not None and runtime.async_session_factory is not None: return try: engine = create_async_db_engine(config=runtime.config) session_factory = create_async_session_factory(engine) except Exception as exc: cls._mark_unavailable(runtime) raise cls._data_source_error(DataSourceInitializationException, runtime, exc) from None runtime.async_engine = engine runtime.async_session_factory = session_factory async def initialize(self, log_enabled: bool = True) -> None: """ 初始化并检查所有数据源的连接状态 :param log_enabled: 是否输出数据源启动日志 :return: None """ self._log_enabled = log_enabled if self._initialized: return names = tuple(self._configs) results = await asyncio.gather(*(self._check_health(name) for name in names), return_exceptions=True) default_name = self._resolve_name() required_failure: tuple[str, BaseException] | None = None for name, result in zip(names, results, strict=True): config = self._configs[name] required = config.db_required or name == default_name healthy = not isinstance(result, BaseException) error_type = None error_code = None if not healthy: error_type, error_code = _error_details(result) if log_enabled: log_context: dict[str, Any] = { 'data_source': name, 'database_type': config.db_type, 'required': required, } if error_type is not None: log_context['error_type'] = error_type log_context['error_code'] = error_code elif healthy: log_context['session_timezone'] = self._runtime(name).session_timezone source_logger = logger.bind(**log_context) if healthy: source_logger.info( f'✅ 数据源 {name} 初始化成功，sessionTimeZone={self._runtime(name).session_timezone}' ) elif required: source_logger.error(f'❌ 必需数据源 {name} 连接检查失败{_error_log_suffix(error_type, error_code)}') else: source_logger.warning( f'⚠️ 非必需数据源 {name} 连接检查失败，应用将降级启动{_error_log_suffix(error_type, error_code)}' ) if not healthy and required and required_failure is None: required_failure = (name, result) if required_failure is not None: name, result = required_failure error_type, error_code = _error_details(result) await self.dispose_all() raise DataSourceInitializationException( name, error_type=error_type, error_code=error_code, ) from None self._initialized = True def get_async_engine(self, name: str | None = None) -> AsyncEngine: """ 获取数据源异步引擎 :param name: 数据源名称 :return: 异步SQLAlchemy Engine """ runtime = self._runtime(name) self._ensure_async_resources(runtime) assert runtime.async_engine is not None return runtime.async_engine def get_sync_engine(self, name: str | None = None) -> Engine: """ 获取数据源同步引擎 :param name: 数据源名称 :return: 同步SQLAlchemy Engine """ runtime = self._runtime(name) if runtime.sync_engine is None: try: runtime.sync_engine = create_sync_db_engine(config=runtime.config) except Exception as exc: raise self._data_source_error(DataSourceInitializationException, runtime, exc) from None return runtime.sync_engine async def _check_health(self, name: str) -> None: """ 检查指定数据源的连接状态 :param name: 数据源名称 :return: None """ runtime = self._runtime(name) async with runtime.health_lock: await self._check_health_locked(runtime) async def _check_health_locked(self, runtime: DataSourceRuntime) -> None: """ 在持有健康检查锁时检查数据源连接状态 :param runtime: 数据源运行时状态 :return: None """ try: self._ensure_async_resources(runtime) assert runtime.async_engine is not None async with runtime.async_engine.begin() as connection: await connection.execute(text('SELECT 1')) timezone_result = await connection.execute(text(_SESSION_TIMEZONE_QUERY[runtime.config.db_type])) session_timezone = str(timezone_result.scalar_one()) if session_timezone not in _UTC_SESSION_TIMEZONES: raise RuntimeError(f'数据库session timezone必须为UTC，实际为{session_timezone}') runtime.session_timezone = session_timezone except Exception as exc: self._mark_unavailable(runtime) raise self._data_source_error(DataSourceUnavailableException, runtime, exc) from None runtime.available = True runtime.last_health_check_at = TimezoneUtil.utc_now() runtime.next_retry_at = None async def _ensure_available(self, runtime: DataSourceRuntime) -> None: """ 确保指定数据源当前可用 :param runtime: 数据源运行时状态 :return: None """ async with runtime.health_lock: if runtime.available: return now = TimezoneUtil.utc_now() if runtime.next_retry_at is not None and now < runtime.next_retry_at: raise DataSourceUnavailableException(runtime.name) await self._check_health_locked(runtime) if self._log_enabled: logger.bind(data_source=runtime.name, session_timezone=runtime.session_timezone).info( f'✅ 数据源 {runtime.name} 连接已恢复，sessionTimeZone={runtime.session_timezone}' ) @asynccontextmanager async def connection(self, name: str | None = None) -> AsyncGenerator[AsyncConnection, None]: """ 创建指定数据源的异步数据库连接事务 :param name: 数据源名称 :return: 异步数据库连接 """ runtime = self._runtime(name) await self._ensure_available(runtime) assert runtime.async_engine is not None try: async with runtime.async_engine.begin() as connection: yield connection except DBAPIError as exc: if not exc.connection_invalidated: raise self._mark_unavailable(runtime) raise self._data_source_error(DataSourceUnavailableException, runtime, exc) from None @asynccontextmanager async def session(self, name: str | None = None, *, log_sql: bool = True) -> AsyncGenerator[AsyncSession, None]: """ 创建指定数据源的异步数据库会话 :param name: 数据源名称 :param log_sql: 是否保留数据源原有的SQL日志设置，False时只过滤此会话的常规SQL日志 :return: 异步数据库会话 """ runtime = self._runtime(name) await self._ensure_available(runtime) factory = runtime.async_session_factory assert factory is not None session_options: dict[str, Any] = {} if not log_sql: assert runtime.async_engine is not None session_options['bind'] = quiet_sql_engine(runtime.async_engine) try: async with factory(**session_options) as current_db: yield current_db except DBAPIError as exc: if not exc.connection_invalidated: raise self._mark_unavailable(runtime) raise self._data_source_error(DataSourceUnavailableException, runtime, exc) from None async def dispose_all(self) -> None: """ 释放所有数据源的同步和异步引擎 :return: None """ runtimes = tuple(self._runtimes.values()) self._runtimes.clear() self._initialized = False for runtime in runtimes: if runtime.sync_engine is not None: try: runtime.sync_engine.dispose() except Exception as exc: error_type, error_code = _error_details(exc) logger.bind( data_source=runtime.name, engine_type='sync', error_type=error_type, error_code=error_code, ).warning(f'⚠️ 数据源 {runtime.name} 同步Engine释放失败{_error_log_suffix(error_type, error_code)}') async_resources = [(runtime, engine) for runtime in runtimes if (engine := runtime.async_engine) is not None] results = await asyncio.gather( *(engine.dispose() for _, engine in async_resources), return_exceptions=True, ) for (runtime, _), result in zip(async_resources, results, strict=True): if isinstance(result, Exception): error_type, error_code = _error_details(result) logger.bind( data_source=runtime.name, engine_type='async', error_type=error_type, error_code=error_code, ).warning(f'⚠️ 数据源 {runtime.name} 异步Engine释放失败{_error_log_suffix(error_type, error_code)}') elif isinstance(result, BaseException): raise result DataSourceRegistry = _DataSourceRegistry() class Base(AsyncAttrs, DeclarativeBase): pass @cache def get_data_source_base(source_name: str) -> type[DeclarativeBase]: """ 获取指定数据源独立且可复用的ORM元数据基类 :param source_name: 数据源名称 :return: ORM元数据基类 """ DataBaseConfig.get_source(source_name) class NamedDataSourceBase(AsyncAttrs, DeclarativeBase): pass NamedDataSourceBase.__name__ = f'{source_name.title().replace("_", "").replace("-", "")}DataSourceBase' return NamedDataSourceBase