harness_zbt/backend/app/infrastructure/mysql/sessions.py

"""SQLAlchemy 数据库连接与会话工具。

Engine 管理数据库连接池；Session 表示一次数据库工作单元，负责查询和提交事务。
"""

from collections.abc import Iterator
from enum import StrEnum

from sqlalchemy import URL, Engine, create_engine
from sqlalchemy.orm import Session, sessionmaker

from app.core.config import Settings


class DatabaseName(StrEnum):
    """三个业务数据库的逻辑名称，使用枚举可避免到处手写字符串。"""

    CORE = "core"
    AGENT = "agent"
    ANALYTICS = "analytics"


def build_mysql_url(settings: Settings, database: DatabaseName) -> URL:
    """根据配置和目标数据库生成 SQLAlchemy 连接地址。"""
    database_names = {
        DatabaseName.CORE: settings.mysql_core_database,
        DatabaseName.AGENT: settings.mysql_agent_database,
        DatabaseName.ANALYTICS: settings.mysql_analytics_database,
    }
    return URL.create(
        drivername="mysql+pymysql",
        username=settings.mysql_user,
        password=settings.mysql_password,
        host=settings.mysql_host,
        port=settings.mysql_port,
        database=database_names[database],
        query={"charset": "utf8mb4"},
    )


def create_mysql_engine(settings: Settings, database: DatabaseName) -> Engine:
    """创建连接池；此时通常还不会立刻执行 SQL。"""
    return create_engine(
        build_mysql_url(settings, database),
        # 借出连接前先检查连接是否有效，避免拿到已被 MySQL 关闭的旧连接。
        pool_pre_ping=True,
        # 定期回收连接，降低 MySQL 服务端超时导致断连的概率。
        pool_recycle=1800,
    )


def create_session_factory(engine: Engine) -> sessionmaker[Session]:
    """创建 Session 工厂；调用这个工厂才能得到具体会话。"""
    return sessionmaker(bind=engine, expire_on_commit=False)


def session_scope(factory: sessionmaker[Session]) -> Iterator[Session]:
    """提供自动提交/回滚的会话上下文，yield 把会话交给调用方使用。"""
    with factory.begin() as session:
        yield session