dc-nogi/bot/storage.py

import sqlite3
import threading
from typing import Optional


SCHEMA = """
CREATE TABLE IF NOT EXISTS targets (
    target_id      TEXT PRIMARY KEY,
    target_type    TEXT NOT NULL CHECK(target_type IN ('guild','group_dm')),
    name           TEXT,
    first_seen_at  TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
    reaction_count INTEGER NOT NULL DEFAULT 0
);

CREATE TABLE IF NOT EXISTS reactions_log (
    id                  INTEGER PRIMARY KEY AUTOINCREMENT,
    message_id          TEXT NOT NULL,
    channel_id          TEXT NOT NULL,
    target_id           TEXT NOT NULL,
    author_id           TEXT NOT NULL,
    emoji               TEXT NOT NULL,
    reacted_at          TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
    message_created_at  TEXT,
    source              TEXT NOT NULL CHECK(source IN ('backlog','live')),
    UNIQUE(message_id, emoji)
);
CREATE INDEX IF NOT EXISTS idx_reactions_log_target ON reactions_log(target_id);

CREATE TABLE IF NOT EXISTS channel_scan_state (
    channel_id        TEXT PRIMARY KEY,
    target_id         TEXT NOT NULL,
    newest_id_seen     TEXT,
    backlog_complete   INTEGER NOT NULL DEFAULT 0,
    last_scanned_at    TEXT
);
"""


class Storage:
    """Thin synchronous sqlite3 wrapper. Callers from async code should
    dispatch through asyncio.to_thread() to avoid blocking the event loop."""

    def __init__(self, db_path: str):
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._lock = threading.Lock()
        with self._lock, self._conn:
            self._conn.execute("PRAGMA journal_mode=WAL")
            self._conn.executescript(SCHEMA)
            self._migrate()

    def _migrate(self) -> None:
        # Older DBs created before message_created_at existed.
        columns = {row["name"] for row in self._conn.execute("PRAGMA table_info(reactions_log)")}
        if "message_created_at" not in columns:
            self._conn.execute("ALTER TABLE reactions_log ADD COLUMN message_created_at TEXT")

    def upsert_target(self, target_id, target_type: str, name: str) -> None:
        with self._lock, self._conn:
            self._conn.execute(
                """
                INSERT INTO targets (target_id, target_type, name)
                VALUES (?, ?, ?)
                ON CONFLICT(target_id) DO UPDATE SET name = excluded.name
                """,
                (str(target_id), target_type, name),
            )

    def has_reacted(self, message_id, emoji: str) -> bool:
        with self._lock:
            cur = self._conn.execute(
                "SELECT 1 FROM reactions_log WHERE message_id = ? AND emoji = ? LIMIT 1",
                (str(message_id), emoji),
            )
            return cur.fetchone() is not None

    def record_reaction(
        self, message_id, channel_id, target_id, author_id, emoji: str, source: str, message_created_at=None
    ) -> bool:
        created_at = message_created_at.strftime("%Y-%m-%dT%H:%M:%fZ") if message_created_at else None
        with self._lock, self._conn:
            cur = self._conn.execute(
                """
                INSERT OR IGNORE INTO reactions_log
                    (message_id, channel_id, target_id, author_id, emoji, source, message_created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (str(message_id), str(channel_id), str(target_id), str(author_id), emoji, source, created_at),
            )
            inserted = cur.rowcount > 0
            if inserted:
                self._conn.execute(
                    "UPDATE targets SET reaction_count = reaction_count + 1 WHERE target_id = ?",
                    (str(target_id),),
                )
            return inserted

    def get_scan_state(self, channel_id) -> Optional[dict]:
        with self._lock:
            cur = self._conn.execute(
                "SELECT * FROM channel_scan_state WHERE channel_id = ?",
                (str(channel_id),),
            )
            row = cur.fetchone()
            return dict(row) if row else None

    def set_scan_state(self, channel_id, target_id, newest_id_seen, backlog_complete: bool) -> None:
        with self._lock, self._conn:
            self._conn.execute(
                """
                INSERT INTO channel_scan_state
                    (channel_id, target_id, newest_id_seen, backlog_complete, last_scanned_at)
                VALUES (?, ?, ?, ?, strftime('%Y-%m-%dT%H:%M:%fZ','now'))
                ON CONFLICT(channel_id) DO UPDATE SET
                    newest_id_seen = excluded.newest_id_seen,
                    backlog_complete = excluded.backlog_complete,
                    last_scanned_at = excluded.last_scanned_at
                """,
                (str(channel_id), str(target_id), str(newest_id_seen) if newest_id_seen else None, int(backlog_complete)),
            )

    def get_all_reactions(self) -> list:
        with self._lock:
            cur = self._conn.execute(
                "SELECT id, message_id, channel_id, target_id, emoji FROM reactions_log ORDER BY id"
            )
            return [dict(row) for row in cur.fetchall()]

    def delete_reaction(self, row_id: int) -> None:
        with self._lock, self._conn:
            row = self._conn.execute(
                "SELECT target_id FROM reactions_log WHERE id = ?", (row_id,)
            ).fetchone()
            self._conn.execute("DELETE FROM reactions_log WHERE id = ?", (row_id,))
            if row:
                self._conn.execute(
                    "UPDATE targets SET reaction_count = MAX(reaction_count - 1, 0) WHERE target_id = ?",
                    (row["target_id"],),
                )

    def get_stats(self) -> list:
        with self._lock:
            cur = self._conn.execute(
                """
                SELECT t.target_id, t.target_type, t.name, t.reaction_count,
                       COUNT(r.id) AS logged_reactions, MAX(r.message_created_at) AS last_reaction_at
                FROM targets t
                LEFT JOIN reactions_log r ON r.target_id = t.target_id
                GROUP BY t.target_id
                ORDER BY t.reaction_count DESC
                """
            )
            return [dict(row) for row in cur.fetchall()]

    def close(self) -> None:
        with self._lock:
            self._conn.close()