e5a725fc91
- database.py: engine/SessionLocal 改为延迟创建，避免 import 时副作用， 并提供 dispose_engine() 支持配置热重载 - dao/uploaded_file_dao.py: 新增 count() 方法 - upload_service.py: list_files() 改用 dao.count() 返回数据库总条数， 修复之前返回当前页条目数导致分页 total 语义错误的问题
71 lines
1.9 KiB
Python
71 lines
1.9 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:
|
||
"""按需建表（幂等）。先导入 models 以注册映射。"""
|
||
from . import models # noqa: F401
|
||
get_engine()
|
||
Base.metadata.create_all(bind=_engine)
|