"""数据库引擎、会话与初始化。 使用 SQLAlchemy 2.0 + SQLite。FTS5 虚拟表通过原生 SQL 创建,并配套触发器 让 OCR/AI 字段更新时自动同步到全文索引。 """ from __future__ import annotations from contextlib import contextmanager from typing import Iterator from sqlalchemy import create_engine, event, text from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker from app.core.config import settings class Base(DeclarativeBase): """全局声明性 Base。""" engine = create_engine( settings.db_url, echo=False, future=True, connect_args={"check_same_thread": False}, ) @event.listens_for(engine, "connect") def _sqlite_pragmas(dbapi_connection, _connection_record): """启用外键、WAL、忙等待等 SQLite 优化项。""" cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA synchronous=NORMAL") cursor.execute("PRAGMA busy_timeout=5000") cursor.close() SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True) def get_session() -> Iterator[Session]: """FastAPI 依赖注入:每个请求一个会话。""" with SessionLocal() as session: yield session @contextmanager def session_scope() -> Iterator[Session]: """常规上下文管理:自动 commit/rollback。""" session = SessionLocal() try: yield session session.commit() except Exception: session.rollback() raise finally: session.close() # FTS5 虚拟表与触发器 SQL(独立维护,便于以后调整字段) _FTS_SCHEMA_SQL = [ """ CREATE VIRTUAL TABLE IF NOT EXISTS screenshots_fts USING fts5( ocr_text, ai_title, ai_summary, ai_suggestion, content='screenshot_meta', content_rowid='screenshot_id', tokenize='unicode61' ); """, """ CREATE TRIGGER IF NOT EXISTS screenshot_meta_ai AFTER INSERT ON screenshot_meta BEGIN INSERT INTO screenshots_fts(rowid, ocr_text, ai_title, ai_summary, ai_suggestion) VALUES (new.screenshot_id, coalesce(new.ocr_text, ''), coalesce(new.ai_title, ''), coalesce(new.ai_summary, ''), coalesce(new.ai_suggestion, '')); END; """, """ CREATE TRIGGER IF NOT EXISTS screenshot_meta_ad AFTER DELETE ON screenshot_meta BEGIN INSERT INTO screenshots_fts(screenshots_fts, rowid, ocr_text, ai_title, ai_summary, ai_suggestion) VALUES('delete', old.screenshot_id, coalesce(old.ocr_text, ''), coalesce(old.ai_title, ''), coalesce(old.ai_summary, ''), coalesce(old.ai_suggestion, '')); END; """, """ CREATE TRIGGER IF NOT EXISTS screenshot_meta_au AFTER UPDATE ON screenshot_meta BEGIN INSERT INTO screenshots_fts(screenshots_fts, rowid, ocr_text, ai_title, ai_summary, ai_suggestion) VALUES('delete', old.screenshot_id, coalesce(old.ocr_text, ''), coalesce(old.ai_title, ''), coalesce(old.ai_summary, ''), coalesce(old.ai_suggestion, '')); INSERT INTO screenshots_fts(rowid, ocr_text, ai_title, ai_summary, ai_suggestion) VALUES (new.screenshot_id, coalesce(new.ocr_text, ''), coalesce(new.ai_title, ''), coalesce(new.ai_summary, ''), coalesce(new.ai_suggestion, '')); END; """, ] def init_db() -> None: """启动时建表并装配 FTS5、灌入默认分类。""" from app.models import register_all # noqa: F401 register_all() Base.metadata.create_all(engine) with engine.begin() as conn: for stmt in _FTS_SCHEMA_SQL: conn.execute(text(stmt)) _migrate_legacy_schema(conn) # 启动期 seed 默认分类(即使首次启动也能在「设置」/筛选页看到分类) from app.services.analyze import ensure_default_categories ensure_default_categories() def _migrate_legacy_schema(conn) -> None: """轻量迁移:旧版本的 screenshots.category_id 没有外键。 SQLite 不支持 ALTER TABLE 加外键,但删除分类时 ON DELETE SET NULL 失效 会导致悬空引用。检测到旧表时,主动用一次性 SQL 清理掉无效引用并打日志, 建议用户用「分类管理」页重建索引。 """ pragma_rows = conn.execute( text("PRAGMA foreign_key_list(screenshots)") ).fetchall() has_cat_fk = any(row[2] == "categories" for row in pragma_rows) if not has_cat_fk: # 清理悬空 category_id,避免列表统计出错 conn.execute( text( "UPDATE screenshots SET category_id = NULL " "WHERE category_id IS NOT NULL " "AND category_id NOT IN (SELECT id FROM categories)" ) )