server-core/app/db.py

# -*- coding: utf-8 -*-
"""SQLAlchemy 引擎 / 会话 / 声明基类。

auth/RBAC 域数据从 JSON 文件迁移到 SQLite，Repository 内部用 SQL 实现，
路由与业务逻辑不变（见 ``repositories.py`` 保留的方法签名）。
"""
from __future__ import annotations

from collections.abc import Generator

from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker

from . import config


class Base(DeclarativeBase):
    """SQLAlchemy 声明式基类。"""


engine = create_engine(
    config.DATABASE_URL,
    connect_args={"check_same_thread": False} if config.DATABASE_URL.startswith("sqlite") else {},
    pool_pre_ping=True,
)

SessionLocal = sessionmaker(
    bind=engine,
    autoflush=False,
    expire_on_commit=False,
)


def init_db() -> None:
    """幂等地创建全部表（create_all 对已存在表无操作）。"""
    # 延迟导入模型以确保表注册到 Base.metadata
    from .infrastructure import models  # noqa: F401

    Base.metadata.create_all(bind=engine)


def create_all(engine) -> None:
    """在指定 engine 上创建全部表（供测试/独立实例使用）。"""
    from .infrastructure import models  # noqa: F401

    Base.metadata.create_all(bind=engine)


def get_session() -> Generator[Session, None, None]:
    """FastAPI 依赖：请求级 SQLAlchemy 会话，请求结束自动关闭。"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def make_engine(db_url: str):
    """按给定 URL 创建独立引擎（测试用 tmp 数据库）。"""
    return create_engine(
        db_url,
        connect_args={"check_same_thread": False} if db_url.startswith("sqlite") else {},
        pool_pre_ping=True,
    )


def make_session_factory(db_url: str) -> sessionmaker:
    return sessionmaker(
        bind=make_engine(db_url),
        autoflush=False,
        expire_on_commit=False,
    )