09682245f5
Base.registry.mappers 在 SQLAlchemy 2.0.x 是 frozenset（装 Mapper 对象），
而非 dict（{table: mapper}），原 .items() 调用会抛
'frozenset' object has no attribute 'items' 导致启动期建表校验失败。
改为统一取 Mapper，用 local_table.name 取表名、columns 取列，兼容 dict 与
frozenset 两种形态。
102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
"""数据库引擎、Session 与 Declarative Base。
|
||
|
||
引擎采用懒初始化：首次访问 `get_engine()` 时才创建连接池，
|
||
避免 import 时的副作用，并支持 `dispose()` 后重新加载配置。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Generator
|
||
|
||
from sqlalchemy import create_engine
|
||
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
||
|
||
from .config import get_settings
|
||
|
||
_engine: create_engine | None = None
|
||
_session_local: sessionmaker | None = None
|
||
|
||
|
||
def get_engine() -> create_engine:
|
||
"""返回全局 Engine 实例；首次调用时懒创建。"""
|
||
global _engine, _session_local
|
||
if _engine is None:
|
||
s = get_settings()
|
||
_engine = create_engine(
|
||
s.db_url(),
|
||
pool_pre_ping=True,
|
||
pool_size=s.database.pool_size,
|
||
pool_recycle=s.database.pool_recycle,
|
||
future=True,
|
||
)
|
||
_session_local = sessionmaker(
|
||
bind=_engine, autoflush=False, autocommit=False, future=True,
|
||
)
|
||
return _engine
|
||
|
||
|
||
def dispose_engine() -> None:
|
||
"""关闭连接池并清除缓存，下次访问时重新创建（配合 reload_settings 使用）。"""
|
||
global _engine, _session_local
|
||
if _engine is not None:
|
||
_engine.dispose()
|
||
_engine = None
|
||
_session_local = None
|
||
|
||
|
||
def get_session_local() -> sessionmaker:
|
||
"""返回全局 SessionLocal；确保 Engine 已初始化。"""
|
||
get_engine()
|
||
return _session_local # type: ignore[return-value]
|
||
|
||
|
||
class Base(DeclarativeBase):
|
||
pass
|
||
|
||
|
||
def get_db() -> Generator[Session, None, None]:
|
||
"""FastAPI 依赖：为每个请求产出一个 Session。"""
|
||
db = get_session_local()()
|
||
try:
|
||
yield db
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
def init_db_schema() -> None:
|
||
"""按需建表（幂等）并校验既有表列与模型一致（fail-fast on schema drift）。
|
||
|
||
先导入 models 注册映射；create_all 用 IF NOT EXISTS 仅补缺失的表；
|
||
随后对每张已存在的表检查模型声明的列是否齐全，缺列即抛 RuntimeError，
|
||
避免运行期才以晦涩的 OperationalError 暴露 schema 漂移。
|
||
"""
|
||
from sqlalchemy import inspect
|
||
|
||
from . import models # noqa: F401
|
||
engine = get_engine()
|
||
Base.metadata.create_all(bind=engine)
|
||
|
||
inspector = inspect(engine)
|
||
missing: list[str] = []
|
||
# Base.registry.mappers 在不同 SQLAlchemy 版本中既可能是 dict（{table: mapper}），
|
||
# 也可能是 frozenset（直接装 Mapper 对象）。统一取 Mapper，用 local_table 取表名、
|
||
# columns 取模型声明的列，兼容两种形态。
|
||
mappers = Base.registry.mappers
|
||
if hasattr(mappers, "values"): # dict 形态
|
||
mapper_iter = mappers.values()
|
||
else: # frozenset 形态（SQLAlchemy 2.0.x）
|
||
mapper_iter = iter(mappers)
|
||
for mapper in mapper_iter:
|
||
table = mapper.local_table.name
|
||
if not inspector.has_table(table):
|
||
continue
|
||
db_cols = {c["name"] for c in inspector.get_columns(table)}
|
||
for model_col in mapper.columns.keys():
|
||
if model_col not in db_cols:
|
||
missing.append(f"{table}.{model_col}")
|
||
if missing:
|
||
raise RuntimeError(
|
||
"数据库 schema 与模型不一致，缺少列: " + ", ".join(missing)
|
||
+ "。请执行 sql/schema.sql 或迁移脚本更新表结构。"
|
||
)
|