8152494fdc
- database.py: echo 改为 False - main.py: 错误响应直接显示异常信息（不再根据 DEBUG 判断） - 机器人重复图片回复增加 OCR 识别内容展示
108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
"""
|
||
数据库连接模块
|
||
支持 PostgreSQL（生产）和 SQLite（本地开发）两种后端
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import AsyncGenerator
|
||
from pathlib import Path
|
||
|
||
from sqlalchemy import text
|
||
from sqlalchemy.ext.asyncio import (
|
||
AsyncSession,
|
||
async_sessionmaker,
|
||
create_async_engine,
|
||
)
|
||
from sqlalchemy.pool import NullPool
|
||
|
||
from app.config import settings
|
||
|
||
# ──────────────────────────── 判断数据库类型 ────────────────────────────
|
||
|
||
DB_URL = settings.DATABASE_URL
|
||
IS_SQLITE = DB_URL.startswith("sqlite")
|
||
|
||
# ──────────────────────────── 异步引擎 ────────────────────────────
|
||
|
||
if IS_SQLITE:
|
||
# SQLite 配置
|
||
sqlite_path = DB_URL.replace("sqlite+aiosqlite:///", "")
|
||
if not sqlite_path.startswith("/"):
|
||
sqlite_path = str(Path(__file__).parent.parent / sqlite_path)
|
||
|
||
engine = create_async_engine(
|
||
f"sqlite+aiosqlite:///{sqlite_path}",
|
||
echo=False,
|
||
)
|
||
else:
|
||
# PostgreSQL 配置
|
||
engine = create_async_engine(
|
||
DB_URL,
|
||
echo=False,
|
||
poolclass=NullPool,
|
||
pool_pre_ping=True,
|
||
)
|
||
|
||
# ──────────────────────────── 会话工厂 ────────────────────────────
|
||
|
||
async_session_factory = async_sessionmaker(
|
||
bind=engine,
|
||
class_=AsyncSession,
|
||
expire_on_commit=False,
|
||
autoflush=False,
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 依赖注入 ────────────────────────────
|
||
|
||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||
"""
|
||
FastAPI 依赖项：获取异步数据库会话。
|
||
"""
|
||
async with async_session_factory() as session:
|
||
try:
|
||
yield session
|
||
await session.commit()
|
||
except Exception:
|
||
await session.rollback()
|
||
raise
|
||
|
||
|
||
# ──────────────────────────── 生命周期辅助 ────────────────────────────
|
||
|
||
async def init_db() -> None:
|
||
"""初始化数据库：建表"""
|
||
from app.models.base import Base
|
||
|
||
async with engine.begin() as conn:
|
||
if IS_SQLITE:
|
||
# 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)
|
||
|
||
if not IS_SQLITE:
|
||
# PostgreSQL 需要创建扩展
|
||
async with engine.begin() as conn:
|
||
try:
|
||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
|
||
except Exception:
|
||
pass
|
||
try:
|
||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS zhparser"))
|
||
except Exception:
|
||
pass
|
||
try:
|
||
await conn.execute(text(
|
||
"CREATE TEXT SEARCH CONFIGURATION IF NOT EXISTS chinese_zh (PARSER = zhparser)"
|
||
))
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
async def close_db() -> None:
|
||
"""关闭数据库引擎"""
|
||
await engine.dispose()
|