xiaoxia-saas/packages/adapters/sqlalchemy_impl/session.py

from __future__ import annotations

from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL, make_url
from sqlalchemy.orm import sessionmaker

from packages.adapters.sqlalchemy_impl.models import Base

SCHEMA_INIT_LOCK_ID = 2026061501
SessionLocal = None


def build_engine(
    database_url: str,
    *,
    pool_size: int = 20,
    max_overflow: int = 40,
    pool_timeout: int = 30,
    pool_recycle: int = 3600,
):
    # SQLite 不支持 QueuePool 的 pool_size/max_overflow/pool_timeout，
    # 传了会在 create_engine 阶段直接 TypeError，这里只对非 SQLite 传连接池参数。
    if _is_sqlite(database_url):
        return create_engine(database_url, pool_recycle=pool_recycle)
    return create_engine(
        database_url,
        pool_size=pool_size,
        max_overflow=max_overflow,
        pool_timeout=pool_timeout,
        pool_recycle=pool_recycle,
    )


def build_session_factory(
    database_url: str,
    *,
    pool_size: int = 20,
    max_overflow: int = 40,
    pool_timeout: int = 30,
    pool_recycle: int = 3600,
):
    engine = build_engine(
        database_url,
        pool_size=pool_size,
        max_overflow=max_overflow,
        pool_timeout=pool_timeout,
        pool_recycle=pool_recycle,
    )
    session_factory = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    global SessionLocal
    SessionLocal = session_factory
    return engine, session_factory


def _is_sqlite(database_url: str) -> bool:
    """检测是否为 SQLite 数据库 URL."""
    return database_url.startswith("sqlite")


def _build_admin_url(database_url: str) -> URL:
    url = make_url(database_url)
    return url.set(database="postgres")


def ensure_database_exists(database_url: str) -> None:
    """确保数据库存在（仅 PostgreSQL 需要，SQLite 自动创建）."""
    if _is_sqlite(database_url):
        return
    target_url = make_url(database_url)
    admin_engine = create_engine(_build_admin_url(database_url), isolation_level="AUTOCOMMIT")
    try:
        with admin_engine.connect() as connection:
            exists = connection.execute(
                text("SELECT 1 FROM pg_database WHERE datname = :database_name"),
                {"database_name": target_url.database},
            ).scalar()
            if exists:
                return
            connection.execute(text(f'CREATE DATABASE "{target_url.database}"'))
    finally:
        admin_engine.dispose()


def initialize_database(engine) -> None:
    """初始化数据库 schema。

    PostgreSQL 使用 advisory lock 防止并发初始化冲突；
    SQLite 直接 create_all（单文件，无并发风险）。
    """
    if _is_sqlite(str(engine.url)):
        Base.metadata.create_all(bind=engine)
        return
    with engine.connect() as connection:
        connection.execute(text("SELECT pg_advisory_lock(:lock_id)"), {"lock_id": SCHEMA_INIT_LOCK_ID})
        try:
            Base.metadata.create_all(bind=connection)
            connection.commit()
        finally:
            connection.execute(
                text("SELECT pg_advisory_unlock(:lock_id)"),
                {"lock_id": SCHEMA_INIT_LOCK_ID},
            )
            connection.commit()