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()


_VIRAL_VIDEO_BACKFILL_COLS = [
    ("storyboard", "JSON"),
    ("generated_copy_text", "TEXT NOT NULL DEFAULT ''"),
]


def _ensure_viral_video_columns(connection) -> None:
    """Idempotently add new columns to viral_video_jobs; create_all will not ALTER existing tables."""
    from sqlalchemy import inspect as _inspect

    try:
        insp = _inspect(connection)
        if not insp.has_table("viral_video_jobs"):
            return
        existing = {c["name"] for c in insp.get_columns("viral_video_jobs")}
    except Exception:
        return
    import logging as _logging

    _log = _logging.getLogger(__name__)
    for col, ddl in _VIRAL_VIDEO_BACKFILL_COLS:
        if col in existing:
            continue
        try:
            connection.execute(text(f"ALTER TABLE viral_video_jobs ADD COLUMN {col} {ddl}"))
            _log.info("added column viral_video_jobs.%s", col)
        except Exception as e:
            _log.warning("add column %s failed: %s", col, e)


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},
            )
            _ensure_viral_video_columns(connection)
            connection.commit()