Initial commit
This commit is contained in:
157
bot/storage.py
Normal file
157
bot/storage.py
Normal file
@@ -0,0 +1,157 @@
|
||||
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')),
|
||||
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")
|
||||
|
||||
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 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()
|
||||
Reference in New Issue
Block a user