SnapAndAnaly/backend/app/core/db.py

"""数据库引擎、会话与初始化。

使用 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)"
            )
        )