# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence. """SQLite persistence for tallies, cooldowns, daily limits and votes. Everything the bot remembers lives here so a restart before the daily report does not lose the day. The original bot kept all of it in module globals. """ from __future__ import annotations import contextlib import sqlite3 from dataclasses import dataclass from typing import TYPE_CHECKING, Self if TYPE_CHECKING: from collections.abc import Iterator from types import TracebackType # Bumped whenever the shape below changes. See Store._apply_schema for why # a version number is the only safe way to notice. SCHEMA_VERSION = 2 # Every table is keyed by group first. The bot watches several groups, and a # tally, a cooldown or a vote that leaked between them would mean a test # message throttling the production group -- so the scope is in the primary # key rather than in the callers' discipline. _SCHEMA = """ CREATE TABLE IF NOT EXISTS message_counts ( group_id TEXT NOT NULL, uuid TEXT NOT NULL, display_name TEXT NOT NULL, count INTEGER NOT NULL DEFAULT 0, PRIMARY KEY (group_id, uuid) ); CREATE TABLE IF NOT EXISTS command_uses ( group_id TEXT NOT NULL, user_uuid TEXT NOT NULL, command TEXT NOT NULL, used_at REAL NOT NULL, PRIMARY KEY (group_id, user_uuid, command) ); CREATE TABLE IF NOT EXISTS daily_uses ( group_id TEXT NOT NULL, user_uuid TEXT NOT NULL, command TEXT NOT NULL, day TEXT NOT NULL, uses INTEGER NOT NULL DEFAULT 0, PRIMARY KEY (group_id, user_uuid, command, day) ); CREATE TABLE IF NOT EXISTS votes ( group_id TEXT NOT NULL, command TEXT NOT NULL, voter_uuid TEXT NOT NULL, voted_at REAL NOT NULL, PRIMARY KEY (group_id, command, voter_uuid) ); """ _TABLES = ("message_counts", "command_uses", "daily_uses", "votes") @dataclass(frozen=True, slots=True) class Tally: """One person's message count for the current reporting period.""" uuid: str display_name: str count: int class Store: """Thin, synchronous SQLite wrapper. All state the bot keeps lives here.""" def __init__(self, path: str) -> None: """Open (or create) the database at path and apply the schema.""" self._conn = sqlite3.connect(path) self._conn.row_factory = sqlite3.Row self._apply_schema() def _apply_schema(self) -> None: """Create the tables, rebuilding them when the file predates them. ``CREATE TABLE IF NOT EXISTS`` is a no-op against a table that already exists with the wrong columns, so a database written before the group_id key would survive startup untouched and then fail on every single INSERT -- a bot that starts clean and answers nothing. PRAGMA user_version is the only thing that notices, because it is the only part of the file that records which shape it holds. The old rows are dropped rather than migrated: they carry no group, and inventing one would silently attribute a whole day's counts to the wrong conversation. """ version = int(self._conn.execute("PRAGMA user_version").fetchone()[0]) if version < SCHEMA_VERSION: for table in _TABLES: self._conn.execute(f"DROP TABLE IF EXISTS {table}") self._conn.executescript(_SCHEMA) self._conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") self._conn.commit() def __enter__(self) -> Self: """Enter a context manager that closes the connection on exit.""" return self def __exit__( self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None, ) -> None: """Close the underlying connection.""" self.close() def close(self) -> None: """Close the underlying connection.""" self._conn.close() def forget_group(self, group_id: str) -> None: """Drop every row belonging to one group. Called when the bot is removed from a group, so that being added back starts a fresh tally and a clean cooldown rather than resurrecting whatever was true whenever it was last in there. """ # Spelled out rather than looped over _TABLES: an f-string here is a # SQL-injection shape even when the input is a module constant, and # the repo does not take lint suppressions. test_forget_group_clears # _every_table is what keeps this list in step with the schema. self._conn.execute("DELETE FROM message_counts WHERE group_id = ?", (group_id,)) self._conn.execute("DELETE FROM command_uses WHERE group_id = ?", (group_id,)) self._conn.execute("DELETE FROM daily_uses WHERE group_id = ?", (group_id,)) self._conn.execute("DELETE FROM votes WHERE group_id = ?", (group_id,)) self._conn.commit() # -- message tallies --------------------------------------------------- def increment(self, group_id: str, uuid: str, display_name: str) -> None: """Add one to a person's tally in one group, creating the row when new.""" self._conn.execute( """ INSERT INTO message_counts (group_id, uuid, display_name, count) VALUES (?, ?, ?, 1) ON CONFLICT(group_id, uuid) DO UPDATE SET count = count + 1, display_name = excluded.display_name """, (group_id, uuid, display_name), ) self._conn.commit() def tallies(self, group_id: str) -> list[Tally]: """Return one group's tallies, highest count first.""" rows = self._conn.execute( "SELECT uuid, display_name, count FROM message_counts " "WHERE group_id = ? ORDER BY count DESC, display_name ASC", (group_id,), ).fetchall() return [Tally(r["uuid"], r["display_name"], r["count"]) for r in rows] def reset_tallies(self, group_id: str) -> None: """Clear one group's tallies, as its daily report does after publishing.""" self._conn.execute("DELETE FROM message_counts WHERE group_id = ?", (group_id,)) self._conn.commit() # -- per-group, per-user, per-command rate limiting -------------------- def record_use(self, group_id: str, user_uuid: str, command: str, now: float, day: str) -> None: """Record one successful use for both cooldown and daily-limit checks.""" self._conn.execute( "INSERT INTO command_uses (group_id, user_uuid, command, used_at) " "VALUES (?, ?, ?, ?) " "ON CONFLICT(group_id, user_uuid, command) DO UPDATE SET used_at = excluded.used_at", (group_id, user_uuid, command, now), ) self._conn.execute( """ INSERT INTO daily_uses (group_id, user_uuid, command, day, uses) VALUES (?, ?, ?, ?, 1) ON CONFLICT(group_id, user_uuid, command, day) DO UPDATE SET uses = uses + 1 """, (group_id, user_uuid, command, day), ) self._conn.commit() def last_used(self, group_id: str, user_uuid: str, command: str) -> float | None: """Return when this user last ran this command in this group, or None.""" row = self._conn.execute( "SELECT used_at FROM command_uses WHERE group_id = ? AND user_uuid = ? AND command = ?", (group_id, user_uuid, command), ).fetchone() return None if row is None else float(row["used_at"]) def uses_today(self, group_id: str, user_uuid: str, command: str, day: str) -> int: """Return how often this user ran this command in this group that day.""" row = self._conn.execute( "SELECT uses FROM daily_uses " "WHERE group_id = ? AND user_uuid = ? AND command = ? AND day = ?", (group_id, user_uuid, command, day), ).fetchone() return 0 if row is None else int(row["uses"]) # -- vote gating ------------------------------------------------------- def add_vote(self, group_id: str, command: str, voter_uuid: str, now: float) -> None: """Record one vote, ignoring a repeat vote from the same person.""" self._conn.execute( "INSERT INTO votes (group_id, command, voter_uuid, voted_at) VALUES (?, ?, ?, ?) " "ON CONFLICT(group_id, command, voter_uuid) DO NOTHING", (group_id, command, voter_uuid, now), ) self._conn.commit() def count_votes(self, group_id: str, command: str, not_before: float) -> int: """Return how many votes for command in this group are still in the window.""" row = self._conn.execute( "SELECT COUNT(*) AS n FROM votes WHERE group_id = ? AND command = ? AND voted_at >= ?", (group_id, command, not_before), ).fetchone() return int(row["n"]) def clear_votes(self, group_id: str, command: str) -> None: """Drop every vote for a command in one group, after it fires or expires.""" self._conn.execute( "DELETE FROM votes WHERE group_id = ? AND command = ?", (group_id, command) ) self._conn.commit() @contextlib.contextmanager def open_store(path: str) -> Iterator[Store]: """Open a Store as a context manager.""" store = Store(path) try: yield store finally: store.close()