from __future__ import annotations from dataclasses import dataclass import time import aiosqlite @dataclass(frozen=True) class PendingVerification: chat_id: int user_id: int challenge_message_id: int expires_at: int class Database: def __init__(self, path: str) -> None: self._path = path self._conn: aiosqlite.Connection | None = None async def connect(self) -> None: self._conn = await aiosqlite.connect(self._path) await self._conn.execute("PRAGMA journal_mode=WAL;") await self._conn.execute("PRAGMA foreign_keys=ON;") await self._conn.execute( """ CREATE TABLE IF NOT EXISTS verified_members ( chat_id INTEGER NOT NULL, user_id INTEGER NOT NULL, verified_at INTEGER NOT NULL, PRIMARY KEY (chat_id, user_id) ); """ ) await self._conn.execute( """ CREATE TABLE IF NOT EXISTS pending_verifications ( chat_id INTEGER NOT NULL, user_id INTEGER NOT NULL, challenge_message_id INTEGER NOT NULL, expires_at INTEGER NOT NULL, PRIMARY KEY (chat_id, user_id) ); """ ) await self._conn.commit() async def close(self) -> None: if self._conn is not None: await self._conn.close() self._conn = None @property def conn(self) -> aiosqlite.Connection: if self._conn is None: raise RuntimeError("Database is not connected") return self._conn async def mark_pending( self, *, chat_id: int, user_id: int, challenge_message_id: int, expires_at: int, ) -> None: await self.conn.execute( """ INSERT INTO pending_verifications (chat_id, user_id, challenge_message_id, expires_at) VALUES (?, ?, ?, ?) ON CONFLICT(chat_id, user_id) DO UPDATE SET challenge_message_id=excluded.challenge_message_id, expires_at=excluded.expires_at """, (chat_id, user_id, challenge_message_id, expires_at), ) await self.conn.commit() async def get_pending(self, *, chat_id: int, user_id: int) -> PendingVerification | None: cur = await self.conn.execute( """ SELECT chat_id, user_id, challenge_message_id, expires_at FROM pending_verifications WHERE chat_id=? AND user_id=? """, (chat_id, user_id), ) row = await cur.fetchone() await cur.close() if not row: return None return PendingVerification( chat_id=int(row[0]), user_id=int(row[1]), challenge_message_id=int(row[2]), expires_at=int(row[3]), ) async def delete_pending(self, *, chat_id: int, user_id: int) -> None: await self.conn.execute( "DELETE FROM pending_verifications WHERE chat_id=? AND user_id=?", (chat_id, user_id), ) await self.conn.commit() async def mark_verified(self, *, chat_id: int, user_id: int) -> None: now = int(time.time()) await self.conn.execute( """ INSERT INTO verified_members (chat_id, user_id, verified_at) VALUES (?, ?, ?) ON CONFLICT(chat_id, user_id) DO UPDATE SET verified_at=excluded.verified_at """, (chat_id, user_id, now), ) await self.conn.execute( "DELETE FROM pending_verifications WHERE chat_id=? AND user_id=?", (chat_id, user_id), ) await self.conn.commit() async def delete_verified(self, *, chat_id: int, user_id: int) -> None: await self.conn.execute( "DELETE FROM verified_members WHERE chat_id=? AND user_id=?", (chat_id, user_id), ) await self.conn.commit() async def is_verified(self, *, chat_id: int, user_id: int) -> bool: cur = await self.conn.execute( """ SELECT 1 FROM verified_members WHERE chat_id=? AND user_id=? """, (chat_id, user_id), ) row = await cur.fetchone() await cur.close() return row is not None async def get_expired_pending(self, *, now: int) -> list[PendingVerification]: cur = await self.conn.execute( """ SELECT chat_id, user_id, challenge_message_id, expires_at FROM pending_verifications WHERE expires_at <= ? ORDER BY expires_at ASC LIMIT 100 """, (now,), ) rows = await cur.fetchall() await cur.close() return [ PendingVerification( chat_id=int(r[0]), user_id=int(r[1]), challenge_message_id=int(r[2]), expires_at=int(r[3]), ) for r in rows ]