catalog_doc/catalogizer/db.py

"""Хранилище каталога: SQLite + полнотекстовый индекс FTS5."""
from __future__ import annotations

import json
import sqlite3
from pathlib import Path

from . import config

SCHEMA = """
PRAGMA journal_mode = WAL;

CREATE TABLE IF NOT EXISTS docs (
    id          INTEGER PRIMARY KEY,
    path        TEXT UNIQUE NOT NULL,   -- абсолютный путь
    rel_path    TEXT NOT NULL,          -- путь от корня библиотеки
    folder      TEXT NOT NULL,          -- раздел (папка) библиотеки
    filename    TEXT NOT NULL,
    ext         TEXT NOT NULL,
    size        INTEGER NOT NULL,
    mtime       REAL NOT NULL,
    title       TEXT,
    author      TEXT,
    year        INTEGER,
    pages       INTEGER,
    lang        TEXT,
    kind        TEXT,
    text_chars  INTEGER DEFAULT 0,
    scanned     INTEGER DEFAULT 0,      -- 1 = нет текстового слоя
    error       TEXT,
    body        TEXT,                   -- извлечённый текст (обрезанный)
    toc         TEXT,                   -- оглавление, JSON-список
    terms       TEXT,                   -- частоты основ, JSON
    description TEXT,
    theses      TEXT,                   -- JSON-список тезисов
    keywords    TEXT,                   -- JSON-список ключевых слов
    stems       TEXT,                   -- основы слов текста для поиска
    indexed_at  TEXT,
    analyzed    INTEGER DEFAULT 0
);

CREATE INDEX IF NOT EXISTS idx_docs_folder ON docs(folder);
CREATE INDEX IF NOT EXISTS idx_docs_ext ON docs(ext);
CREATE INDEX IF NOT EXISTS idx_docs_analyzed ON docs(analyzed);

CREATE VIRTUAL TABLE IF NOT EXISTS docs_fts USING fts5(
    title, keywords, description, body, stems, path,
    tokenize = 'unicode61 remove_diacritics 2'
);

CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT);

-- Ключевые слова, добавленные вручную. Хранятся по пути файла, а не по id:
-- так они переживают перестроение каталога и повторный разбор папки.
CREATE TABLE IF NOT EXISTS user_keywords (
    path     TEXT PRIMARY KEY,
    keywords TEXT NOT NULL,             -- JSON-список
    changed  TEXT
);

-- Счётчик попыток разбора. Нужен, чтобы файл, на котором парсер аварийно
-- завершает весь процесс, не блокировал обход при следующем запуске.
CREATE TABLE IF NOT EXISTS attempts (path TEXT PRIMARY KEY, n INTEGER NOT NULL);
"""


#: Колонки, добавленные после первых версий базы. Старые каталоги
#: дополняются на месте, чтобы их не пришлось строить заново.
MIGRATIONS = (("docs", "stems", "TEXT"),)


def connect(db_path: Path | None = None,
            create: bool = True) -> sqlite3.Connection:
    """Подключение к базе каталога.

    @param create Создавать базу, если её ещё нет. Окно читает каталоги
        с ``create=False``: просмотр не должен плодить пустые базы рядом
        с программой.
    """
    path = Path(db_path or config.DB_PATH)
    if not create and not path.exists():
        raise FileNotFoundError("каталог ещё не построен: %s" % path)
    path.parent.mkdir(parents=True, exist_ok=True)
    con = sqlite3.connect(path, timeout=60)
    con.row_factory = sqlite3.Row
    con.executescript(SCHEMA)
    for table, column, decl in MIGRATIONS:
        have = {r["name"] for r in con.execute("PRAGMA table_info(%s)" % table)}
        if column not in have:
            con.execute("ALTER TABLE %s ADD COLUMN %s %s" % (table, column, decl))
    con.commit()
    return con


def get_state(con: sqlite3.Connection) -> dict[str, tuple[float, int]]:
    """Соответствие путь -> (mtime, size) для инкрементального обхода."""
    return {r["path"]: (r["mtime"], r["size"])
            for r in con.execute("SELECT path, mtime, size FROM docs")}


def upsert_doc(con: sqlite3.Connection, rec: dict) -> None:
    cols = ("path", "rel_path", "folder", "filename", "ext", "size", "mtime",
            "title", "author", "year", "pages", "lang", "kind", "text_chars",
            "scanned", "error", "body", "toc", "terms", "indexed_at")
    values = [rec.get(c) for c in cols]
    con.execute(
        "INSERT INTO docs (%s, analyzed) VALUES (%s, 0) "
        "ON CONFLICT(path) DO UPDATE SET %s, analyzed = 0"
        % (", ".join(cols),
           ", ".join("?" * len(cols)),
           ", ".join("%s = excluded.%s" % (c, c) for c in cols[1:])),
        values,
    )


def save_summary(con: sqlite3.Connection, doc_id: int, description: str,
                 theses: list[str], keywords: list[str], stems: str) -> None:
    con.execute(
        "UPDATE docs SET description = ?, theses = ?, keywords = ?, stems = ?, "
        "analyzed = 1 WHERE id = ?",
        (description, json.dumps(theses, ensure_ascii=False),
         json.dumps(keywords, ensure_ascii=False), stems, doc_id))
    reindex_doc(con, doc_id)


def reindex_doc(con: sqlite3.Connection, doc_id: int) -> None:
    """Перестраивает строку полнотекстового индекса для одного документа.

    Вызывается и после анализа, и после правки ключевых слов вручную —
    иначе добавленное слово не находилось бы поиском.
    """
    row = con.execute(
        "SELECT d.title, d.body, d.path, d.description, d.keywords, d.stems,"
        "       u.keywords AS mine"
        "  FROM docs d LEFT JOIN user_keywords u ON u.path = d.path"
        " WHERE d.id = ?", (doc_id,)).fetchone()
    if row is None:
        return
    words = _as_list(row["keywords"]) + _as_list(row["mine"])
    con.execute("DELETE FROM docs_fts WHERE rowid = ?", (doc_id,))
    con.execute(
        "INSERT INTO docs_fts (rowid, title, keywords, description, body, stems,"
        " path) VALUES (?, ?, ?, ?, ?, ?, ?)",
        (doc_id, row["title"] or "", " ".join(words), row["description"] or "",
         (row["body"] or "")[:20_000], row["stems"] or "", row["path"]))


def _as_list(value: str | None) -> list[str]:
    try:
        return list(json.loads(value or "[]"))
    except (TypeError, ValueError):
        return []


def get_user_keywords(con: sqlite3.Connection, path: str) -> list[str]:
    """Ключевые слова, добавленные к документу вручную."""
    row = con.execute("SELECT keywords FROM user_keywords WHERE path = ?",
                      (path,)).fetchone()
    return _as_list(row["keywords"]) if row else []


def set_user_keywords(con: sqlite3.Connection, path: str,
                      keywords: list[str]) -> list[str]:
    """Заменяет свои ключевые слова документа и обновляет поисковый индекс.

    @param keywords Новый список; пустой список удаляет запись.
    @return Сохранённый список без повторов и пустых значений.
    """
    from datetime import datetime

    clean: list[str] = []
    for word in keywords:
        word = " ".join(str(word).split())
        if word and word.lower() not in {w.lower() for w in clean}:
            clean.append(word)
    if clean:
        con.execute(
            "INSERT INTO user_keywords (path, keywords, changed)"
            " VALUES (?, ?, ?) ON CONFLICT(path) DO UPDATE SET"
            " keywords = excluded.keywords, changed = excluded.changed",
            (path, json.dumps(clean, ensure_ascii=False),
             datetime.now().isoformat(timespec="seconds")))
    else:
        con.execute("DELETE FROM user_keywords WHERE path = ?", (path,))
    row = con.execute("SELECT id FROM docs WHERE path = ?", (path,)).fetchone()
    if row:
        reindex_doc(con, row["id"])
    con.commit()
    return clean


def all_user_keywords(con: sqlite3.Connection) -> list[str]:
    """Все свои ключевые слова каталога — для подсказок при вводе."""
    seen: dict[str, None] = {}
    for row in con.execute("SELECT keywords FROM user_keywords"):
        for word in _as_list(row["keywords"]):
            seen.setdefault(word, None)
    return sorted(seen, key=str.lower)


def purge_missing(con: sqlite3.Connection, alive: set[str]) -> int:
    """Удаляет из базы записи о файлах, которых больше нет на диске."""
    gone = [r["id"] for r in con.execute("SELECT id, path FROM docs")
            if r["path"] not in alive]
    for i in gone:
        con.execute("DELETE FROM docs WHERE id = ?", (i,))
        con.execute("DELETE FROM docs_fts WHERE rowid = ?", (i,))
    return len(gone)


def set_meta(con: sqlite3.Connection, key: str, value: str) -> None:
    con.execute("INSERT INTO meta (key, value) VALUES (?, ?) "
                "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
                (key, value))


def get_meta(con: sqlite3.Connection, key: str, default: str = "") -> str:
    row = con.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone()
    return row["value"] if row else default


def bump_attempt(con: sqlite3.Connection, path: str) -> None:
    con.execute("INSERT INTO attempts (path, n) VALUES (?, 1) "
                "ON CONFLICT(path) DO UPDATE SET n = n + 1", (path,))


def clear_attempt(con: sqlite3.Connection, path: str) -> None:
    con.execute("DELETE FROM attempts WHERE path = ?", (path,))


def failed_paths(con: sqlite3.Connection, max_tries: int = 3) -> set[str]:
    """Файлы, разбор которых уже несколько раз обрывал процесс."""
    return {r["path"] for r in
            con.execute("SELECT path FROM attempts WHERE n >= ?", (max_tries,))}