baoxiang/backend/app/core/database.py

"""
数据库连接和会话管理
"""
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool
from .config import settings

# 创建数据库引擎，支持MySQL连接池
db_settings = settings.database

if "sqlite" in db_settings.DATABASE_URL:
    # SQLite 配置
    engine = create_engine(
        db_settings.DATABASE_URL,
        echo=False,
        future=True,
        connect_args={"check_same_thread": False}  # SQLite 特定配置
    )
else:
    # MySQL 配置
    engine = create_engine(
        db_settings.DATABASE_URL,
        # 连接池配置
        poolclass=QueuePool,
        pool_size=db_settings.DB_POOL_SIZE,
        max_overflow=db_settings.DB_MAX_OVERFLOW,
        pool_timeout=db_settings.DB_POOL_TIMEOUT,
        pool_recycle=db_settings.DB_POOL_RECYCLE,
        echo=False,
        future=True,
        # 设置连接时使用的字符集
        connect_args={"charset": "utf8mb4", "collation": "utf8mb4_unicode_ci"}
    )

# 如果是MySQL，启用严格模式
if "mysql" in db_settings.DATABASE_URL:
    @event.listens_for(engine, "connect")
    def set_sqlite_pragma(dbapi_connection, connection_record):
        """
        为MySQL连接设置参数（如果使用MySQL）
        """
        # 设置字符集为utf8mb4
        try:
            cursor = dbapi_connection.cursor()
            cursor.execute("SET NAMES utf8mb4 COLLATE utf8mb4_unicode_ci")
            cursor.close()
        except Exception:
            pass


# 创建会话
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db():
    """
    获取数据库会话
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def create_tables():
    """
    创建所有表
    """
    from ..models.base import Base
    from ..models import user, game
    Base.metadata.create_all(bind=engine)


def drop_tables():
    """
    删除所有表（仅用于开发环境）
    """
    from ..models.base import Base
    from ..models import user, game
    Base.metadata.drop_all(bind=engine)