Rain-Bus
85c4cfb9e8
- Add image_extractor, pdf_downloader, pi_client, trends services - Add shared utils module - Refactor summarizer, embedder, routes for cleaner separation - Update tests to match new service structure
74 lines
1.8 KiB
Python
74 lines
1.8 KiB
Python
"""数据库引擎、会话工厂、初始化。"""
|
|
|
|
from sqlalchemy import event, create_engine, text
|
|
from sqlalchemy.orm import DeclarativeBase, sessionmaker
|
|
|
|
from app.config import settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
# ── FTS5 和索引 DDL（与 ORM 模型分开管理）───────────────────────────────
|
|
|
|
FTS5_CREATE_SQL = """
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS papers_fts USING fts5(
|
|
title_en,
|
|
title_zh,
|
|
abstract,
|
|
authors,
|
|
tags,
|
|
summary_text,
|
|
tokenize='unicode61'
|
|
);
|
|
"""
|
|
|
|
FTS5_TRIGGER_INDEX = """
|
|
-- partial index for task_locks running
|
|
CREATE UNIQUE INDEX IF NOT EXISTS uq_task_locks_running
|
|
ON task_locks(task, lock_key) WHERE status = 'running';
|
|
"""
|
|
|
|
|
|
def _make_engine():
|
|
"""创建 SQLite 引擎，启用 foreign_keys。"""
|
|
engine = create_engine(
|
|
settings.DATABASE_URL,
|
|
echo=settings.APP_DEBUG,
|
|
connect_args={"check_same_thread": False},
|
|
)
|
|
|
|
@event.listens_for(engine, "connect")
|
|
def _set_sqlite_pragma(dbapi_connection, _connection_record):
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.close()
|
|
|
|
return engine
|
|
|
|
|
|
engine = _make_engine()
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
|
|
|
|
|
def get_db():
|
|
"""FastAPI 依赖注入：获取数据库会话。"""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def init_db(engine):
|
|
"""创建所有 ORM 表 + FTS5 虚拟表。"""
|
|
from app.models import Base # noqa: F811 — 避免循环导入，延迟导入
|
|
|
|
Base.metadata.create_all(engine)
|
|
with engine.connect() as conn:
|
|
conn.execute(text(FTS5_CREATE_SQL))
|
|
conn.execute(text(FTS5_TRIGGER_INDEX))
|
|
conn.commit()
|