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.
"""

from __future__ import annotations

import contextlib
import sqlite3
from typing import TYPE_CHECKING, Self

from signal_bot.schema import EPHEMERAL_TABLES, SCHEMA, SCHEMA_VERSION
from signal_bot.stats.events import EventStore

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


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)

    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:
            # EPHEMERAL_TABLES, not TABLES. Rebuilding a cooldown table is
            # free; rebuilding message_events would throw away the entire
            # statistical history, which nothing can recompute.
            for table in EPHEMERAL_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 one group's *ephemeral* rows: 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 _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 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()