81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
"""
|
||
数据库连接管理
|
||
|
||
SQLite 数据库，支持异步操作。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from typing import AsyncGenerator
|
||
|
||
from sqlalchemy import text
|
||
from sqlalchemy.ext.asyncio import (
|
||
AsyncSession,
|
||
async_sessionmaker,
|
||
create_async_engine,
|
||
)
|
||
|
||
from app.config import settings
|
||
|
||
# ──────────────────────────── 数据库引擎 ────────────────────────────
|
||
|
||
DB_URL = settings.DATABASE_URL
|
||
|
||
# 解析 SQLite 文件路径
|
||
sqlite_path = DB_URL.replace("sqlite+aiosqlite:///", "")
|
||
if not sqlite_path.startswith("/"):
|
||
sqlite_path = str(Path(__file__).parent.parent / sqlite_path)
|
||
|
||
# 确保数据目录存在
|
||
Path(sqlite_path).parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
engine = create_async_engine(
|
||
f"sqlite+aiosqlite:///{sqlite_path}",
|
||
echo=False,
|
||
pool_size=1,
|
||
max_overflow=0,
|
||
)
|
||
|
||
# ──────────────────────────── 会话工厂 ────────────────────────────
|
||
|
||
async_session_factory = async_sessionmaker(
|
||
bind=engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False,
|
||
autoflush=False,
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 依赖注入 ────────────────────────────
|
||
|
||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||
"""
|
||
FastAPI 依赖项：获取异步数据库会话。
|
||
请求正常结束后自动 commit，异常时自动 rollback。
|
||
"""
|
||
async with async_session_factory() as session:
|
||
try:
|
||
yield session
|
||
await session.commit()
|
||
except Exception:
|
||
await session.rollback()
|
||
raise
|
||
|
||
|
||
async def close_db() -> None:
|
||
"""关闭数据库连接池"""
|
||
await engine.dispose()
|
||
|
||
|
||
async def init_db() -> None:
|
||
"""初始化数据库：建表"""
|
||
from app.models.base import Base
|
||
|
||
async with engine.begin() as conn:
|
||
# SQLite 启用 WAL 模式和外键约束
|
||
await conn.execute(text("PRAGMA journal_mode=WAL"))
|
||
await conn.execute(text("PRAGMA foreign_keys=ON"))
|
||
# 创建所有表
|
||
await conn.run_sync(Base.metadata.create_all)
|