76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
"""SQLAlchemy 2.x 同步引擎 + 会话管理（MVP: SQLite）。
|
|
|
|
切换 PostgreSQL 时（扩展接口 1）：
|
|
1. DATABASE_URL 改为 postgresql+asyncpg://...
|
|
2. 引擎换 create_async_engine + async_sessionmaker
|
|
3. get_session 改 async generator + yield
|
|
4. 业务层 Repository 调用加 await
|
|
"""
|
|
|
|
from collections.abc import Generator
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from app.core.config import get_settings
|
|
|
|
_engine = None
|
|
_session_factory: sessionmaker[Session] | None = None
|
|
|
|
|
|
def get_engine():
|
|
global _engine
|
|
if _engine is None:
|
|
settings = get_settings()
|
|
_engine = create_engine(
|
|
settings.database_url,
|
|
pool_pre_ping=True,
|
|
echo=False,
|
|
# SQLite 专属：启用 WAL 模式（并发读 + 写串行化）
|
|
connect_args={"check_same_thread": False} if "sqlite" in settings.database_url else {},
|
|
)
|
|
# SQLite: 启用 WAL 模式与外键约束
|
|
if "sqlite" in settings.database_url:
|
|
from sqlalchemy import event, text
|
|
|
|
@event.listens_for(_engine, "connect")
|
|
def _set_sqlite_pragma(dbapi_conn, _): # type: ignore[no-untyped-def]
|
|
cursor = dbapi_conn.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.close()
|
|
|
|
return _engine
|
|
|
|
|
|
def get_session_factory() -> sessionmaker[Session]:
|
|
global _session_factory
|
|
if _session_factory is None:
|
|
_session_factory = sessionmaker(
|
|
get_engine(),
|
|
class_=Session,
|
|
expire_on_commit=False,
|
|
)
|
|
return _session_factory
|
|
|
|
|
|
def get_session() -> Generator[Session, None, None]:
|
|
"""FastAPI 依赖：请求级会话。"""
|
|
factory = get_session_factory()
|
|
session = factory()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
def dispose_engine() -> None:
|
|
global _engine, _session_factory
|
|
if _engine is not None:
|
|
_engine.dispose()
|
|
_engine = None
|
|
_session_factory = None |