news/news/store.py

"""SQLite schema and repository helpers.

Text ids and RFC3339 timestamps throughout, so the database stays readable with
the `sqlite3` CLI. Vectors are stored as raw float32 blobs.
"""

from __future__ import annotations

import sqlite3
from collections.abc import Iterable, Sequence
from pathlib import Path

import numpy as np

SCHEMA = """
CREATE TABLE IF NOT EXISTS sources (
  id                 TEXT PRIMARY KEY,
  name               TEXT NOT NULL,
  url                TEXT NOT NULL UNIQUE,
  kind               TEXT NOT NULL,              -- rss | hn | reddit
  -- Editorial-independence bucket. Importance counts DISTINCT buckets, not
  -- feeds: ten outlets reprinting one wire story is one bucket, not ten.
  bucket             TEXT NOT NULL,
  weight             REAL NOT NULL DEFAULT 1.0,
  -- 1 for a publisher's front-page/top-stories feed. Appearing there is a free
  -- editorial-prominence signal; the old system dropped it by unsubscribing.
  prominence         INTEGER NOT NULL DEFAULT 0,
  enabled            INTEGER NOT NULL DEFAULT 1,
  poll_interval_secs INTEGER NOT NULL DEFAULT 900,
  etag               TEXT,
  last_modified      TEXT,
  last_polled_at     TEXT,
  last_status        TEXT
);

CREATE TABLE IF NOT EXISTS items (
  id            TEXT PRIMARY KEY,
  source_id     TEXT NOT NULL REFERENCES sources(id),
  title         TEXT NOT NULL,
  summary       TEXT NOT NULL DEFAULT '',
  url           TEXT NOT NULL,
  canonical_url TEXT NOT NULL,
  published_at  TEXT NOT NULL,
  first_seen_at TEXT NOT NULL,
  content_hash  TEXT NOT NULL,
  -- Aggregator-native signal (HN points, Reddit score). Re-polled over the
  -- item's life: an HN post at 3 points is noise, at 300 it is news.
  ext_score     INTEGER,
  ext_comments  INTEGER,
  ext_updated_at TEXT,
  UNIQUE(source_id, canonical_url)
);
CREATE INDEX IF NOT EXISTS items_published ON items(published_at);
CREATE INDEX IF NOT EXISTS items_hash      ON items(content_hash);

-- Keyed by content hash, not item id, so identical text across feeds is
-- embedded exactly once. Embedding is the slow step (~11 items/sec).
CREATE TABLE IF NOT EXISTS embeddings (
  content_hash TEXT PRIMARY KEY,
  dim          INTEGER NOT NULL,
  vec          BLOB NOT NULL
);

CREATE TABLE IF NOT EXISTS clusters (
  id            TEXT PRIMARY KEY,
  centroid      BLOB NOT NULL,
  size          INTEGER NOT NULL DEFAULT 0,
  first_seen_at TEXT NOT NULL,
  last_seen_at  TEXT NOT NULL,
  lead_item_id  TEXT REFERENCES items(id),
  headline      TEXT
);
CREATE INDEX IF NOT EXISTS clusters_last_seen ON clusters(last_seen_at);

CREATE TABLE IF NOT EXISTS cluster_members (
  cluster_id TEXT NOT NULL REFERENCES clusters(id) ON DELETE CASCADE,
  item_id    TEXT NOT NULL REFERENCES items(id),
  added_at   TEXT NOT NULL,
  PRIMARY KEY (cluster_id, item_id)
);
CREATE INDEX IF NOT EXISTS cluster_members_item ON cluster_members(item_id);

-- The suppression and budget ledger. Every send is recorded, including dry
-- runs, so `simulate` and the live router share one notion of "already sent".
CREATE TABLE IF NOT EXISTS notification_log (
  id                INTEGER PRIMARY KEY AUTOINCREMENT,
  cluster_id        TEXT NOT NULL,
  lane              TEXT NOT NULL,
  score             REAL,
  origins_at_notify INTEGER,
  dry_run           INTEGER NOT NULL DEFAULT 0,
  created_at        TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS notification_log_created ON notification_log(created_at);
CREATE INDEX IF NOT EXISTS notification_log_cluster ON notification_log(cluster_id);

-- Append-only: every click is a training event, never an update in place.
CREATE TABLE IF NOT EXISTS feedback (
  id         INTEGER PRIMARY KEY AUTOINCREMENT,
  cluster_id TEXT NOT NULL,
  kind       TEXT NOT NULL,                     -- interested | not_interested
  source     TEXT NOT NULL DEFAULT 'web',       -- web | ntfy
  created_at TEXT NOT NULL
);
"""


def connect(path: str | Path) -> sqlite3.Connection:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    con = sqlite3.connect(path, isolation_level=None)
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=WAL")
    con.execute("PRAGMA foreign_keys=ON")
    con.executescript(SCHEMA)
    return con


def to_blob(vec: np.ndarray) -> bytes:
    return np.asarray(vec, dtype=np.float32).tobytes()


def from_blob(blob: bytes) -> np.ndarray:
    return np.frombuffer(blob, dtype=np.float32)


def upsert_source(con: sqlite3.Connection, **kw) -> None:
    """Config is authoritative for everything except the polling state columns."""
    con.execute(
        """INSERT INTO sources (id,name,url,kind,bucket,weight,prominence,
                                enabled,poll_interval_secs)
           VALUES (:id,:name,:url,:kind,:bucket,:weight,:prominence,
                   :enabled,:poll_interval_secs)
           ON CONFLICT(url) DO UPDATE SET
             name=excluded.name, kind=excluded.kind, bucket=excluded.bucket,
             weight=excluded.weight, prominence=excluded.prominence,
             enabled=excluded.enabled,
             poll_interval_secs=excluded.poll_interval_secs""",
        kw,
    )


def insert_items(con: sqlite3.Connection, rows: Iterable[dict]) -> int:
    """Returns the number of genuinely NEW items.

    Counted by row count, not `total_changes`: the ON CONFLICT branch updates
    aggregator scores on every re-poll, and counting those as new would make an
    HN source look like it produced hundreds of fresh items an hour.
    """
    before = con.execute("SELECT COUNT(*) AS n FROM items").fetchone()["n"]
    con.executemany(
        """INSERT INTO items (id,source_id,title,summary,url,canonical_url,
                              published_at,first_seen_at,content_hash,
                              ext_score,ext_comments,ext_updated_at)
           VALUES (:id,:source_id,:title,:summary,:url,:canonical_url,
                   :published_at,:first_seen_at,:content_hash,
                   :ext_score,:ext_comments,:ext_updated_at)
           ON CONFLICT(source_id,canonical_url) DO UPDATE SET
             ext_score=excluded.ext_score,
             ext_comments=excluded.ext_comments,
             ext_updated_at=excluded.ext_updated_at
           WHERE excluded.ext_score IS NOT NULL""",
        list(rows),
    )
    return con.execute("SELECT COUNT(*) AS n FROM items").fetchone()["n"] - before


def cached_vectors(con: sqlite3.Connection, hashes: Sequence[str]) -> dict[str, np.ndarray]:
    out: dict[str, np.ndarray] = {}
    for i in range(0, len(hashes), 500):
        chunk = hashes[i : i + 500]
        q = ",".join("?" * len(chunk))
        for row in con.execute(
            f"SELECT content_hash, vec FROM embeddings WHERE content_hash IN ({q})", chunk
        ):
            out[row["content_hash"]] = from_blob(row["vec"])
    return out


def store_vectors(con: sqlite3.Connection, pairs: Iterable[tuple[str, np.ndarray]]) -> None:
    con.executemany(
        "INSERT OR REPLACE INTO embeddings (content_hash,dim,vec) VALUES (?,?,?)",
        [(h, len(v), to_blob(v)) for h, v in pairs],
    )