mirror of
https://github.com/kuhyx/signal-bot.git
synced 2026-11-10 18:17:02 +01:00
Krzysztof kuhy Rudnicki
0c4b259c17
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
238 lines
9.3 KiB
Python
238 lines
9.3 KiB
Python
# 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()
|