signal-bot/signal_bot/store.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""SQLite persistence for cooldowns, daily limits, votes and group state.

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.

The message history itself is reached through ``store.events`` -- see
``signal_bot.stats.events``. It is kept apart because it is the one thing here
that is append-only and never discarded. ``store.suggestions`` is the second
such thing: what the group asked for and what it mistyped, which no amount of
recomputation can bring back once dropped.
"""

from __future__ import annotations

import contextlib
import logging
import sqlite3
from typing import TYPE_CHECKING, Self

from signal_bot.analytics.uses import CommandEventStore
from signal_bot.schema import REBUILT_ON_BUMP, RETAINED_COLUMNS, SCHEMA, SCHEMA_VERSION
from signal_bot.schema_commands import COMMAND_SCHEMA
from signal_bot.stats.events import EventStore
from signal_bot.suggestions import SuggestionStore

if TYPE_CHECKING:
    from collections.abc import Iterator
    from types import TracebackType

logger = logging.getLogger(__name__)


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()
        self.events = EventStore(self._conn)
        self.suggestions = SuggestionStore(self._conn)
        self.commands = CommandEventStore(self._conn)

    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:
            # REBUILT_ON_BUMP, not TABLES. Rebuilding a vote in progress is
            # free; rebuilding message_events would throw away the entire
            # statistical history, which nothing can recompute. command_uses
            # and daily_uses used to be in here and are deliberately not any
            # more -- see the comment on the tuple in schema.py.
            for table in REBUILT_ON_BUMP:
                self._conn.execute(f"DROP TABLE IF EXISTS {table}")
            self._drop_reshaped()
        self._conn.executescript(SCHEMA)
        # A separate string, not a separate schema: one database, one
        # user_version. Split only so both files stay under the cap.
        self._conn.executescript(COMMAND_SCHEMA)
        self._conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
        self._conn.commit()

    def _drop_reshaped(self) -> None:
        """Drop a retained table whose columns no longer match the schema.

        command_uses and daily_uses survive a version bump so that the use
        history survives with them, which means a bump can no longer repair
        them. A database written before the group key still holds the old
        two-column command_uses, and CREATE TABLE IF NOT EXISTS is a no-op
        against it -- so without this the bot boots clean and then fails on
        every INSERT, which is the exact failure the version check exists to
        prevent.

        Keeping rows is worth a lot; keeping them at the price of a bot that
        answers nothing is not. A mismatch is logged, because losing a
        cooldown silently is fine and losing history silently is not.
        """
        for table, expected in RETAINED_COLUMNS.items():
            # PRAGMA takes no parameters, and the name is a module constant
            # walked here rather than anything a message could reach.
            found = tuple(
                str(row["name"]) for row in self._conn.execute(f"PRAGMA table_info({table})")
            )
            if found and found != expected:
                logger.warning(
                    "%s has columns %s, expected %s; dropping it to stay writable",
                    table,
                    found,
                    expected,
                )
                self._conn.execute(f"DROP TABLE IF EXISTS {table}")

    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 one group's live scaffolding: cooldowns, limits, votes, meta.

        Called when the bot is removed from a group, so that being added back
        starts with a clean cooldown rather than resurrecting whatever was
        true whenever it was last in there.

        What it deliberately no longer touches is the message history. A group
        listing that transiently omits members -- the exact failure the
        comment in discovery._permitted warns about -- would otherwise delete
        years of statistics on the strength of one bad HTTP response. History
        is permanent; only the live scaffolding is disposable.
        """
        # Spelled out rather than looped over CLEARED_ON_LEAVE: 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_cleared_table is what keeps this
        # list in step with that tuple.
        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.execute("DELETE FROM group_meta WHERE group_id = ?", (group_id,))
        self._conn.commit()

    # -- per-group key/value state -----------------------------------------

    def meta(self, group_id: str, key: str) -> str | None:
        """Return a stored value for one group, or None if never set."""
        row = self._conn.execute(
            "SELECT value FROM group_meta WHERE group_id = ? AND key = ?",
            (group_id, key),
        ).fetchone()
        return None if row is None else str(row["value"])

    def set_meta(self, group_id: str, key: str, value: str) -> None:
        """Store a value for one group, replacing any previous one."""
        self._conn.execute(
            "INSERT INTO group_meta (group_id, key, value) VALUES (?, ?, ?) "
            "ON CONFLICT(group_id, key) DO UPDATE SET value = excluded.value",
            (group_id, key, value),
        )
        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()