76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
"""数据库会话与引擎管理。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncGenerator
|
|
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy import event
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from .core.config import settings
|
|
|
|
DATABASE_URL = settings.database_url
|
|
|
|
|
|
def _is_sqlite() -> bool:
|
|
return DATABASE_URL.lower().startswith("sqlite")
|
|
|
|
|
|
engine_kwargs = {
|
|
"echo": settings.debug,
|
|
"future": True,
|
|
"pool_pre_ping": not _is_sqlite(),
|
|
}
|
|
if _is_sqlite():
|
|
engine_kwargs["connect_args"] = {
|
|
"timeout": 30,
|
|
"check_same_thread": False,
|
|
}
|
|
engine_kwargs["poolclass"] = NullPool
|
|
if not _is_sqlite():
|
|
engine_kwargs.update(
|
|
{
|
|
"pool_size": max(5, int(settings.db_pool_size)),
|
|
"max_overflow": max(0, int(settings.db_max_overflow)),
|
|
"pool_timeout": max(5, int(settings.db_pool_timeout)),
|
|
"pool_recycle": max(60, int(settings.db_pool_recycle)),
|
|
# 强制数据库会话时区为北京时间，确保 now()/CURRENT_TIMESTAMP 为 Asia/Shanghai。
|
|
"connect_args": {"server_settings": {"timezone": "Asia/Shanghai"}},
|
|
}
|
|
)
|
|
|
|
engine: AsyncEngine = create_async_engine(
|
|
DATABASE_URL,
|
|
**engine_kwargs,
|
|
)
|
|
|
|
if _is_sqlite():
|
|
@event.listens_for(engine.sync_engine, "connect")
|
|
def _configure_sqlite_pragmas(dbapi_connection, _connection_record) -> None: # pragma: no cover
|
|
cursor = dbapi_connection.cursor()
|
|
try:
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.execute("PRAGMA synchronous=NORMAL")
|
|
cursor.execute("PRAGMA busy_timeout=30000")
|
|
finally:
|
|
cursor.close()
|
|
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
bind=engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autoflush=False,
|
|
)
|
|
|
|
|
|
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
|
"""FastAPI 依赖：提供数据库会话。"""
|
|
async with AsyncSessionLocal() as session:
|
|
yield session
|