kami_apple_exchage/backend/app/core/database.py

"""
数据库连接管理
支持SQLite和PostgreSQL的同步/异步操作
"""

import asyncio
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Generator

from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import Session
from sqlalchemy.pool import StaticPool

from app.core.config import get_settings
from app.core.log import get_logger

settings = get_settings()
logger = get_logger(__name__)


class DatabaseManager:
    """数据库管理器"""

    def __init__(self):
        self.async_engine = None
        self.async_session_factory = None
        self._setup_engines()

    def _setup_engines(self):
        """设置数据库引擎"""
        logger.info(f"正在配置数据库连接: {settings.DATABASE_URL}")

        # 同步引擎配置

        # 异步引擎配置 - 确保使用asyncpg驱动
        async_url: str = settings.DATABASE_URL

        async_url = async_url.replace("sqlite:///", "sqlite+aiosqlite:///")

        if async_url.startswith("sqlite"):
            self.async_engine = create_async_engine(
                async_url,
                poolclass=StaticPool,
                connect_args={"check_same_thread": False},
                echo=settings.DEBUG,
            )
        else:
            self.async_engine = create_async_engine(
                async_url,
                pool_size=settings.DATABASE_POOL_SIZE,
                max_overflow=settings.DATABASE_MAX_OVERFLOW,
                pool_timeout=settings.DATABASE_TIMEOUT,
                pool_pre_ping=True,
                echo=settings.DEBUG,
            )

        # 创建会话工厂
        # 创建会话工厂
        self.sync_session_factory = None  # 暂时禁用同步会话

        self.async_session_factory = async_sessionmaker(
            bind=self.async_engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )

        logger.info("✅ 数据库引擎配置完成")

    def get_engine(self) -> AsyncEngine:
        """获取数据库连接字符串"""
        if self.async_engine:
            return self.async_engine
        else:
            raise RuntimeError("数据库未初始化，请先调用 initialize() 方法")

    def get_sync_session(self) -> Session:
        """获取同步数据库会话"""
        if self.sync_session_factory is None:
            raise RuntimeError("数据库未初始化，请先调用 initialize() 方法")
        return self.sync_session_factory()

    @asynccontextmanager
    async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取异步数据库会话"""
        if self.async_session_factory is None:
            raise RuntimeError("数据库未初始化，请先调用 initialize() 方法")

        async with self.async_session_factory() as session:
            try:
                yield session
            except Exception:
                await session.rollback()
                raise
            finally:
                await session.close()

    async def close(self):
        """关闭数据库连接"""
        try:
            # 先标记为已关闭状态，防止新的连接
            engines_to_close = []

            if self.async_engine:
                engines_to_close.append(("async", self.async_engine))
                self.async_engine = None

            # 重置会话工厂
            self.async_session_factory = None
            self.sync_session_factory = None

            # 安全关闭引擎
            for engine_type, engine in engines_to_close:
                try:
                    if engine_type == "async":
                        # 异步引擎需要在当前事件循环中关闭
                        await self._safe_dispose_async_engine(engine)
                    else:
                        # 同步引擎可以直接关闭
                        engine.dispose()
                    logger.debug(f"✅ {engine_type}数据库引擎已关闭")
                except Exception as e:
                    logger.warning(f"关闭{engine_type}数据库引擎时出现警告: {e}")
                    # 继续处理其他引擎，不中断整个关闭流程
                    continue

            logger.info("✅ 数据库连接已关闭")
        except Exception as e:
            logger.error(f"关闭数据库连接时发生错误: {e}")
            # 确保引擎变量被重置，即使关闭失败
            self.async_engine = None
            self.async_session_factory = None

    async def _safe_dispose_async_engine(self, engine):
        """安全关闭异步数据库引擎"""
        try:
            # 确保在正确的事件循环中运行
            current_loop = asyncio.get_running_loop()

            # 设置较短的超时时间，避免挂起
            await asyncio.wait_for(engine.dispose(), timeout=10.0)

        except asyncio.TimeoutError:
            logger.warning("数据库引擎关闭超时，强制关闭")
            # 强制关闭，不等待
            try:
                # 获取连接池并强制关闭所有连接
                if hasattr(engine, "_pool"):
                    pool = engine._pool
                    if pool is not None:
                        # 强制关闭连接池
                        await pool.dispose()
            except Exception as force_close_error:
                logger.debug(f"强制关闭连接池时的错误（可忽略）: {force_close_error}")
        except RuntimeError as re:
            if "different loop" in str(re):
                logger.warning("检测到事件循环冲突，使用降级关闭策略")
                # 降级策略：不等待，直接标记为已关闭
                try:
                    # 尝试同步方式强制关闭
                    if hasattr(engine, "sync_engine"):
                        engine.sync_engine.dispose()
                except Exception:
                    pass
            else:
                raise
        except Exception as e:
            logger.warning(f"异步引擎关闭时出现错误: {e}")
            # 记录错误但不中断流程


# 全局数据库管理器实例
db_manager = DatabaseManager()


# 依赖注入函数
def get_sync_db() -> Generator[Session, None, None]:
    """获取同步数据库会话 - 用于依赖注入"""
    db = db_manager.get_sync_session()
    try:
        yield db
    finally:
        db.close()


async def get_async_db() -> AsyncGenerator[AsyncSession, None]:
    """获取异步数据库会话 - 用于依赖注入"""
    async with db_manager.get_async_session() as session:
        yield session


# 生命周期管理函数
async def init_database():
    """初始化数据库"""
    try:
        # 测试数据库连接
        # 测试数据库连接
        from sqlalchemy import text

        async with db_manager.get_async_session() as session:
            await session.execute(text(text="SELECT 1"))
        logger.info("✅ 数据库连接测试成功")
    except Exception as e:
        logger.error(f"❌ 数据库连接失败: {e}")
        raise


async def close_database():
    """关闭数据库连接"""
    try:
        logger.info("开始关闭数据库连接...")
        await db_manager.close()
        logger.info("数据库连接关闭完成")
    except Exception as e:
        logger.error(f"关闭数据库时发生错误: {e}")
        # 确保即使出错也不影响其他关闭流程
        pass


# 兼容性别名
get_db = get_sync_db  # 向后兼容