初始化

app/core/database.py

@@ -0,0 +1,79 @@
+"""双数据源：PostgreSQL（源库） + MySQL（业务库）"""
+
+from typing import Optional
+
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+from sqlalchemy.orm import DeclarativeBase
+
+from app.config import settings
+from app.core.logger import log
+
+# ──────────── 内部变量 ────────────
+_pg_engine: Optional[AsyncEngine] = None
+_pg_session_factory: Optional[async_sessionmaker[AsyncSession]] = None
+
+_mysql_engine: Optional[AsyncEngine] = None
+_mysql_session_factory: Optional[async_sessionmaker[AsyncSession]] = None
+
+
+class PgBase(DeclarativeBase):
+    """PostgreSQL ORM 声明基类"""
+    pass
+
+
+class MysqlBase(DeclarativeBase):
+    """MySQL ORM 声明基类"""
+    pass
+
+
+async def init_db() -> None:
+    """初始化双数据源"""
+    global _pg_engine, _pg_session_factory, _mysql_engine, _mysql_session_factory
+
+    _pg_engine = create_async_engine(
+        settings.pg_url,
+        pool_size=settings.pg_pool_size,
+        max_overflow=settings.pg_max_overflow,
+        pool_recycle=3600,
+        echo=False,
+    )
+    _pg_session_factory = async_sessionmaker(_pg_engine, expire_on_commit=False)
+
+    _mysql_engine = create_async_engine(
+        settings.mysql_url,
+        pool_size=settings.mysql_pool_size,
+        max_overflow=settings.mysql_max_overflow,
+        pool_recycle=3600,
+        echo=False,
+    )
+    _mysql_session_factory = async_sessionmaker(_mysql_engine, expire_on_commit=False)
+
+    log.info("双数据源初始化完成: PG={}, MySQL={}", settings.pg_host, settings.db_host)
+
+
+async def close_db() -> None:
+    """关闭双数据源"""
+    if _pg_engine:
+        await _pg_engine.dispose()
+    if _mysql_engine:
+        await _mysql_engine.dispose()
+    log.info("双数据源已关闭")
+
+
+def PgSession() -> AsyncSession:
+    """获取 PostgreSQL 异步会话（用作 async with PgSession() as session）"""
+    if _pg_session_factory is None:
+        raise RuntimeError("数据库未初始化，请先调用 init_db()")
+    return _pg_session_factory()
+
+
+def MysqlSession() -> AsyncSession:
+    """获取 MySQL 异步会话（用作 async with MysqlSession() as session）"""
+    if _mysql_session_factory is None:
+        raise RuntimeError("数据库未初始化，请先调用 init_db()")
+    return _mysql_session_factory()