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','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") # Older DBs created before 'dm' targets (1:1 DMs) existed have a CHECK # constraint that only allows 'guild'/'group_dm' — rebuild the table # with the wider constraint, preserving all existing rows. targets_sql = self._conn.execute( "SELECT sql FROM sqlite_master WHERE type='table' AND name='targets'" ).fetchone() if targets_sql and "'dm'" not in targets_sql["sql"]: self._conn.executescript( """ ALTER TABLE targets RENAME TO targets_old; CREATE TABLE targets ( target_id TEXT PRIMARY KEY, target_type TEXT NOT NULL CHECK(target_type IN ('guild','group_dm','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 ); INSERT INTO targets SELECT * FROM targets_old; DROP TABLE targets_old; """ ) 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 prune_targets(self, keep_target_ids: list) -> int: """Remove targets not present in the latest discovery run (e.g. a server that's no longer mutual, or a stale entry from an earlier bug/test run). reactions_log/channel_scan_state rows are left alone for history/undo.""" keep = [str(t) for t in keep_target_ids] with self._lock, self._conn: placeholders = ",".join("?" * len(keep)) if keep else None if placeholders: cur = self._conn.execute( f"DELETE FROM targets WHERE target_id NOT IN ({placeholders})", keep ) else: cur = self._conn.execute("DELETE FROM targets") return cur.rowcount 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()