bindery/bindery/store.py

"""SQLite storage for tracked stories, chapter revisions, images and replacement rules.

The store, not the epub, is the source of truth for a story. Every epub is built
from the stored chapters, so rules can change and chapters can disappear from the
site without losing text.
"""

import json
import re
import sqlite3
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from datetime import datetime
from pathlib import Path
from typing import Any

SCHEMA = """
CREATE TABLE IF NOT EXISTS stories (
    id INTEGER PRIMARY KEY,
    url TEXT NOT NULL UNIQUE,
    site TEXT NOT NULL,
    title TEXT NOT NULL DEFAULT '',
    author TEXT NOT NULL DEFAULT '',
    calibre_id INTEGER,
    calibre_uuid TEXT,
    metahtml TEXT,
    locked_fields TEXT NOT NULL DEFAULT '[]',
    pinned TEXT NOT NULL DEFAULT '{}',
    site_values TEXT NOT NULL DEFAULT '{}',
    auto_update INTEGER NOT NULL DEFAULT 1,
    status TEXT NOT NULL DEFAULT '',
    message TEXT NOT NULL DEFAULT '',
    ingested_at TEXT,
    last_checked TEXT,
    last_changed TEXT,
    created_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS chapters (
    id INTEGER PRIMARY KEY,
    story_id INTEGER NOT NULL REFERENCES stories(id) ON DELETE CASCADE,
    position INTEGER NOT NULL,
    url TEXT NOT NULL,
    title TEXT NOT NULL,
    on_site INTEGER NOT NULL DEFAULT 1,
    included INTEGER NOT NULL DEFAULT 1,
    current_rev INTEGER NOT NULL DEFAULT 1,
    added_at TEXT NOT NULL,
    changed_at TEXT,
    UNIQUE (story_id, url)
);
CREATE TABLE IF NOT EXISTS revisions (
    chapter_id INTEGER NOT NULL REFERENCES chapters(id) ON DELETE CASCADE,
    rev INTEGER NOT NULL,
    html TEXT NOT NULL,
    text_hash TEXT NOT NULL,
    source TEXT NOT NULL,
    fetched_at TEXT NOT NULL,
    PRIMARY KEY (chapter_id, rev)
);
CREATE TABLE IF NOT EXISTS images (
    story_id INTEGER NOT NULL REFERENCES stories(id) ON DELETE CASCADE,
    url TEXT NOT NULL,
    ext TEXT NOT NULL,
    mime TEXT NOT NULL,
    data BLOB NOT NULL,
    is_cover INTEGER NOT NULL DEFAULT 0,
    PRIMARY KEY (story_id, url, is_cover)
);
CREATE TABLE IF NOT EXISTS rules (
    id INTEGER PRIMARY KEY,
    story_id INTEGER REFERENCES stories(id) ON DELETE CASCADE,
    position INTEGER NOT NULL,
    pattern TEXT NOT NULL,
    replacement TEXT NOT NULL DEFAULT '',
    flags TEXT NOT NULL DEFAULT '',
    enabled INTEGER NOT NULL DEFAULT 1,
    note TEXT NOT NULL DEFAULT ''
);
CREATE TABLE IF NOT EXISTS settings (
    story_id INTEGER NOT NULL,  -- 0 for the global settings
    name TEXT NOT NULL,
    value TEXT NOT NULL,
    PRIMARY KEY (story_id, name)
);
CREATE INDEX IF NOT EXISTS chapters_story ON chapters (story_id, position);
CREATE INDEX IF NOT EXISTS rules_story ON rules (story_id, position);
"""

_JSON_COLUMNS = ("locked_fields", "pinned", "site_values")
_STORY_COLUMNS = {
    "url",
    "site",
    "title",
    "author",
    "calibre_id",
    "calibre_uuid",
    "metahtml",
    "locked_fields",
    "pinned",
    "site_values",
    "auto_update",
    "status",
    "message",
    "ingested_at",
    "last_checked",
    "last_changed",
}


def now() -> str:
    return datetime.now().isoformat(timespec="seconds")


def _story(row: sqlite3.Row | None) -> dict[str, Any] | None:
    if row is None:
        return None
    d = dict(row)
    for col in _JSON_COLUMNS:
        if col in d:
            d[col] = json.loads(d[col] or ("[]" if col == "locked_fields" else "{}"))
    return d


_ARTICLE = re.compile(r"^(?:the|an?)\s+", re.I)


def sort_title(title: str | None) -> str:
    """The key that sorts a story list by title, without a leading "The", "A" or "An"."""
    return _ARTICLE.sub("", (title or "").strip()).casefold()


class Store:
    def __init__(self, path: Path) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with self._connect() as conn:
            conn.execute("PRAGMA journal_mode=WAL")
            conn.executescript(SCHEMA)

    @contextmanager
    def _connect(self) -> Iterator[sqlite3.Connection]:
        conn = sqlite3.connect(self.path, timeout=30)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys=ON")
        try:
            with conn:
                yield conn
        finally:
            conn.close()

    def backup(self, path: Path) -> None:
        """Write a consistent copy of the database, also while jobs write to it."""
        with self._connect() as conn:
            target = sqlite3.connect(path)
            try:
                conn.backup(target)
            finally:
                target.close()

    # Stories

    def create_story(
        self,
        url: str,
        site: str,
        title: str,
        author: str,
        locked_fields: list[str],
        calibre_id: int | None = None,
        calibre_uuid: str | None = None,
    ) -> int:
        with self._connect() as conn:
            cur = conn.execute(
                "INSERT INTO stories (url, site, title, author, calibre_id, calibre_uuid, locked_fields, created_at) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (url, site, title, author, calibre_id, calibre_uuid, json.dumps(locked_fields), now()),
            )
            return int(cur.lastrowid)

    def get_story(self, story_id: int) -> dict[str, Any] | None:
        with self._connect() as conn:
            return _story(conn.execute("SELECT * FROM stories WHERE id = ?", (story_id,)).fetchone())

    def get_story_by_url(self, url: str) -> dict[str, Any] | None:
        with self._connect() as conn:
            return _story(conn.execute("SELECT * FROM stories WHERE url = ?", (url,)).fetchone())

    def get_story_by_calibre_id(self, calibre_id: int) -> dict[str, Any] | None:
        with self._connect() as conn:
            return _story(conn.execute("SELECT * FROM stories WHERE calibre_id = ?", (calibre_id,)).fetchone())

    def list_stories(self) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute("""
                SELECT s.id, s.url, s.site, s.title, s.author, s.calibre_id, s.auto_update, s.status,
                       s.message, s.last_checked, s.last_changed, s.ingested_at, s.locked_fields,
                       (SELECT COUNT(*) FROM chapters c WHERE c.story_id = s.id AND c.included = 1) AS chapter_count,
                       (SELECT COUNT(*) FROM chapters c WHERE c.story_id = s.id AND c.included = 1
                                                          AND c.on_site = 0) AS offsite_count,
                       (SELECT COUNT(*) FROM chapters c WHERE c.story_id = s.id AND c.included = 1
                        AND (SELECT COUNT(*) FROM revisions r WHERE r.chapter_id = c.id) > 1) AS revised_count
                FROM stories s""").fetchall()
            return [_story(r) for r in sorted(rows, key=lambda r: sort_title(r["title"]))]

    def update_story(self, story_id: int, **fields: Any) -> None:
        if not fields:
            return
        unknown = set(fields) - _STORY_COLUMNS
        if unknown:
            raise ValueError(f"Unknown story columns: {sorted(unknown)}")
        values = [json.dumps(v) if k in _JSON_COLUMNS else v for k, v in fields.items()]
        assignments = ", ".join(f"{k} = ?" for k in fields)
        with self._connect() as conn:
            conn.execute(f"UPDATE stories SET {assignments} WHERE id = ?", (*values, story_id))  # noqa: S608 (checked columns)

    def delete_story(self, story_id: int) -> None:
        with self._connect() as conn:
            conn.execute("DELETE FROM stories WHERE id = ?", (story_id,))
            conn.execute("DELETE FROM settings WHERE story_id = ?", (story_id,))

    # Chapters

    def list_chapters(self, story_id: int) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT c.*, (SELECT COUNT(*) FROM revisions r WHERE r.chapter_id = c.id) AS revisions
                FROM chapters c WHERE c.story_id = ? ORDER BY c.position""",
                (story_id,),
            ).fetchall()
            return [dict(r) for r in rows]

    def chapters_for_build(self, story_id: int) -> list[dict[str, Any]]:
        """Included chapters in book order, with the HTML of their current revision."""
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT c.id, c.url, c.title, c.position, r.html
                FROM chapters c JOIN revisions r ON r.chapter_id = c.id AND r.rev = c.current_rev
                WHERE c.story_id = ? AND c.included = 1 ORDER BY c.position""",
                (story_id,),
            ).fetchall()
            return [dict(r) for r in rows]

    def replace_revision(self, chapter_id: int, rev: int, html: str, text_hash: str) -> None:
        with self._connect() as conn:
            conn.execute(
                "UPDATE revisions SET html = ?, text_hash = ? WHERE chapter_id = ? AND rev = ?",
                (html, text_hash, chapter_id, rev),
            )

    def get_chapter(self, chapter_id: int) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute("SELECT * FROM chapters WHERE id = ?", (chapter_id,)).fetchone()
            return dict(row) if row else None

    @staticmethod
    def _insert_chapter(
        conn: sqlite3.Connection, story_id: int, chapter: dict[str, str], source: str, position: int
    ) -> int:
        """Store a chapter with its first revision. The chapter needs url, title, html and text_hash."""
        cur = conn.execute(
            "INSERT INTO chapters (story_id, position, url, title, added_at) VALUES (?, ?, ?, ?, ?)",
            (story_id, position, chapter["url"], chapter["title"], now()),
        )
        chapter_id = int(cur.lastrowid)
        conn.execute(
            "INSERT INTO revisions (chapter_id, rev, html, text_hash, source, fetched_at) VALUES (?, 1, ?, ?, ?, ?)",
            (chapter_id, chapter["html"], chapter["text_hash"], source, now()),
        )
        return chapter_id

    def add_chapter(
        self, story_id: int, url: str, title: str, html: str, text_hash: str, source: str, position: int
    ) -> int:
        """Add one chapter at a position. The jobs use add_chapters, which also places the chapters."""
        chapter = {"url": url, "title": title, "html": html, "text_hash": text_hash}
        with self._connect() as conn:
            return self._insert_chapter(conn, story_id, chapter, source, position)

    def add_chapters(
        self,
        story_id: int,
        chapters: list[dict[str, str]],
        source: str,
        order: list[str] | Callable[[list[str]], list[str]],
    ) -> list[int]:
        """Add chapters and put the story in the order of `order`, a list of chapter URLs, in one transaction.

        Each chapter needs url, title, html and text_hash. Chapters that `order`
        leaves out go after the others and keep their order.

        `order` can be a function. It gets the URLs of the stored chapters in book
        order, and returns the list. It runs in the transaction, so it sees a move or
        a delete that the user made a moment before.
        """
        with self._connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            if callable(order):
                stored = conn.execute("SELECT url FROM chapters WHERE story_id = ? ORDER BY position, id", (story_id,))
                order = order([r["url"] for r in stored])
            ids = [self._insert_chapter(conn, story_id, c, source, 0) for c in chapters]
            rows = conn.execute(
                "SELECT id, url FROM chapters WHERE story_id = ? ORDER BY position, id", (story_id,)
            ).fetchall()
            by_url = {r["url"]: r["id"] for r in rows}
            placed = [by_url[u] for u in dict.fromkeys(order) if u in by_url]
            seen = set(placed)
            rest = [r["id"] for r in rows if r["id"] not in seen]
            conn.executemany("UPDATE chapters SET position = ? WHERE id = ?", list(enumerate(placed + rest)))
            return ids

    def add_revision(self, chapter_id: int, html: str, text_hash: str, source: str) -> int:
        """Store a new revision and make it the current one."""
        with self._connect() as conn:
            rev = conn.execute(
                "SELECT COALESCE(MAX(rev), 0) + 1 FROM revisions WHERE chapter_id = ?", (chapter_id,)
            ).fetchone()[0]
            conn.execute(
                "INSERT INTO revisions (chapter_id, rev, html, text_hash, source, fetched_at) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                (chapter_id, rev, html, text_hash, source, now()),
            )
            conn.execute("UPDATE chapters SET current_rev = ?, changed_at = ? WHERE id = ?", (rev, now(), chapter_id))
            return int(rev)

    def update_chapter(self, chapter_id: int, **fields: Any) -> None:
        if not fields:
            return
        unknown = set(fields) - {"position", "url", "title", "on_site", "included", "current_rev", "changed_at"}
        if unknown:
            raise ValueError(f"Unknown chapter columns: {sorted(unknown)}")
        assignments = ", ".join(f"{k} = ?" for k in fields)
        with self._connect() as conn:
            conn.execute(
                f"UPDATE chapters SET {assignments} WHERE id = ?",  # noqa: S608 (checked columns)
                (*fields.values(), chapter_id),
            )

    def reorder_chapters(self, story_id: int, arrange: Callable[[list[int]], list[int]]) -> None:
        """Give the chapters of a story a new order.

        `arrange` gets the chapter ids in book order, and returns them in the new
        order. The read and the write are one transaction, so a job that adds a
        chapter at the same time cannot come between them.
        """
        with self._connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            rows = conn.execute("SELECT id FROM chapters WHERE story_id = ? ORDER BY position, id", (story_id,))
            order = arrange([r["id"] for r in rows])
            conn.executemany(
                "UPDATE chapters SET position = ? WHERE id = ? AND story_id = ?",
                [(i, cid, story_id) for i, cid in enumerate(order)],
            )

    def delete_chapters(self, story_id: int, chapter_ids: list[int]) -> None:
        """Delete chapters of the story and their revisions."""
        with self._connect() as conn:
            conn.executemany(
                "DELETE FROM chapters WHERE id = ? AND story_id = ?", [(cid, story_id) for cid in chapter_ids]
            )

    def list_revisions(self, chapter_id: int) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT rev, text_hash, source, fetched_at, LENGTH(html) AS size FROM revisions "
                "WHERE chapter_id = ? ORDER BY rev",
                (chapter_id,),
            ).fetchall()
            return [dict(r) for r in rows]

    def delete_revision(self, chapter_id: int, rev: int) -> None:
        with self._connect() as conn:
            conn.execute("DELETE FROM revisions WHERE chapter_id = ? AND rev = ?", (chapter_id, rev))

    def revision_html(self, chapter_id: int, rev: int) -> str | None:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT html FROM revisions WHERE chapter_id = ? AND rev = ?", (chapter_id, rev)
            ).fetchone()
            return row["html"] if row else None

    # Images

    def put_image(self, story_id: int, url: str, ext: str, mime: str, data: bytes, is_cover: bool = False) -> None:
        with self._connect() as conn:
            if is_cover:
                conn.execute("DELETE FROM images WHERE story_id = ? AND is_cover = 1", (story_id,))
            conn.execute(
                "INSERT OR REPLACE INTO images (story_id, url, ext, mime, data, is_cover) VALUES (?, ?, ?, ?, ?, ?)",
                (story_id, url, ext, mime, data, int(is_cover)),
            )

    def image_urls(self, story_id: int) -> set:
        with self._connect() as conn:
            rows = conn.execute("SELECT url FROM images WHERE story_id = ? AND is_cover = 0", (story_id,))
            return {r["url"] for r in rows}

    def image_sizes(self, story_id: int) -> list[list[Any]]:
        """The URL, the cover mark and the size of each image, without its data."""
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT url, is_cover, LENGTH(data) FROM images WHERE story_id = ? ORDER BY url, is_cover", (story_id,)
            )
            return [list(r) for r in rows]

    def images(self, story_id: int) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT url, ext, mime, data FROM images WHERE story_id = ? AND is_cover = 0", (story_id,)
            ).fetchall()
            return [dict(r) for r in rows]

    def cover(self, story_id: int) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT url, ext, mime, data FROM images WHERE story_id = ? AND is_cover = 1", (story_id,)
            ).fetchone()
            return dict(row) if row else None

    # Replacement rules. story_id None holds the global rules.

    def rules(self, story_id: int | None) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT pattern, replacement, flags, enabled, note FROM rules WHERE story_id IS ? ORDER BY position",
                (story_id,),
            ).fetchall()
            return [dict(r, enabled=bool(r["enabled"])) for r in rows]

    def set_rules(self, story_id: int | None, rules: list[dict[str, Any]]) -> None:
        with self._connect() as conn:
            conn.execute("DELETE FROM rules WHERE story_id IS ?", (story_id,))
            conn.executemany(
                "INSERT INTO rules (story_id, position, pattern, replacement, flags, enabled, note) "
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                [
                    (
                        story_id,
                        i,
                        r["pattern"],
                        r.get("replacement", ""),
                        r.get("flags", ""),
                        int(bool(r.get("enabled", True))),
                        r.get("note", ""),
                    )
                    for i, r in enumerate(rules)
                ],
            )

    # Settings. Story 0 holds the global settings.

    def get_setting(self, story_id: int | None, name: str) -> Any:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT value FROM settings WHERE story_id = ? AND name = ?", (story_id or 0, name)
            ).fetchone()
            return json.loads(row["value"]) if row else None

    def story_settings(self, name: str) -> dict[int, Any]:
        """The value of one setting for each story that has it."""
        with self._connect() as conn:
            rows = conn.execute("SELECT story_id, value FROM settings WHERE name = ? AND story_id != 0", (name,))
            return {r["story_id"]: json.loads(r["value"]) for r in rows}

    def set_setting(self, story_id: int | None, name: str, value: Any) -> None:
        """Store a setting. None removes it."""
        with self._connect() as conn:
            if value is None:
                conn.execute("DELETE FROM settings WHERE story_id = ? AND name = ?", (story_id or 0, name))
            else:
                conn.execute(
                    "INSERT OR REPLACE INTO settings (story_id, name, value) VALUES (?, ?, ?)",
                    (story_id or 0, name, json.dumps(value)),
                )