from __future__ import annotations from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Optional import aiosqlite SCHEMA = """ CREATE TABLE IF NOT EXISTS clients ( telegram_id INTEGER PRIMARY KEY, username TEXT, first_name TEXT, last_name TEXT, topic_id INTEGER UNIQUE, language TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); CREATE INDEX IF NOT EXISTS idx_clients_topic_id ON clients(topic_id); CREATE TABLE IF NOT EXISTS message_links ( id INTEGER PRIMARY KEY AUTOINCREMENT, client_telegram_id INTEGER NOT NULL, client_message_id INTEGER NOT NULL, manager_message_id INTEGER NOT NULL, created_at TEXT NOT NULL, UNIQUE(client_telegram_id, client_message_id), UNIQUE(manager_message_id) ); CREATE INDEX IF NOT EXISTS idx_message_links_manager ON message_links(manager_message_id); """ @dataclass class Client: telegram_id: int username: Optional[str] first_name: Optional[str] last_name: Optional[str] topic_id: Optional[int] language: Optional[str] created_at: str updated_at: str @dataclass class MessageLink: client_telegram_id: int client_message_id: int manager_message_id: int def _utcnow() -> str: return datetime.now(timezone.utc).isoformat() class ClientRepository: def __init__(self, db_path: Path) -> None: self._db_path = db_path self._conn: Optional[aiosqlite.Connection] = None async def connect(self) -> None: self._db_path.parent.mkdir(parents=True, exist_ok=True) self._conn = await aiosqlite.connect(self._db_path) self._conn.row_factory = aiosqlite.Row await self._conn.execute("PRAGMA journal_mode=WAL") await self._conn.execute("PRAGMA busy_timeout=5000") await self._conn.executescript(SCHEMA) await self._ensure_columns() await self._conn.commit() async def _ensure_columns(self) -> None: conn = self._require_conn() async with conn.execute("PRAGMA table_info(clients)") as cursor: rows = await cursor.fetchall() columns = {row[1] for row in rows} if "language" not in columns: await conn.execute("ALTER TABLE clients ADD COLUMN language TEXT") async def close(self) -> None: if self._conn is not None: await self._conn.close() self._conn = None def _require_conn(self) -> aiosqlite.Connection: if self._conn is None: raise RuntimeError("Database is not connected") return self._conn @staticmethod def _row_to_client(row: aiosqlite.Row) -> Client: return Client( telegram_id=row["telegram_id"], username=row["username"], first_name=row["first_name"], last_name=row["last_name"], topic_id=row["topic_id"], language=row["language"], created_at=row["created_at"], updated_at=row["updated_at"], ) async def upsert_client( self, telegram_id: int, username: Optional[str], first_name: Optional[str], last_name: Optional[str], ) -> Client: conn = self._require_conn() now = _utcnow() await conn.execute( """ INSERT INTO clients ( telegram_id, username, first_name, last_name, topic_id, language, created_at, updated_at ) VALUES (?, ?, ?, ?, NULL, NULL, ?, ?) ON CONFLICT(telegram_id) DO UPDATE SET username = excluded.username, first_name = excluded.first_name, last_name = excluded.last_name, updated_at = excluded.updated_at """, (telegram_id, username, first_name, last_name, now, now), ) await conn.commit() client = await self.get_by_telegram_id(telegram_id) if client is None: raise RuntimeError(f"Failed to upsert client {telegram_id}") return client async def get_by_telegram_id(self, telegram_id: int) -> Optional[Client]: conn = self._require_conn() async with conn.execute( "SELECT * FROM clients WHERE telegram_id = ?", (telegram_id,), ) as cursor: row = await cursor.fetchone() return self._row_to_client(row) if row else None async def get_by_topic_id(self, topic_id: int) -> Optional[Client]: conn = self._require_conn() async with conn.execute( "SELECT * FROM clients WHERE topic_id = ?", (topic_id,), ) as cursor: row = await cursor.fetchone() return self._row_to_client(row) if row else None async def set_topic_id(self, telegram_id: int, topic_id: int) -> None: conn = self._require_conn() await conn.execute( """ UPDATE clients SET topic_id = ?, updated_at = ? WHERE telegram_id = ? """, (topic_id, _utcnow(), telegram_id), ) await conn.commit() async def clear_topic_id(self, telegram_id: int) -> None: conn = self._require_conn() await conn.execute( """ UPDATE clients SET topic_id = NULL, updated_at = ? WHERE telegram_id = ? """, (_utcnow(), telegram_id), ) await conn.commit() async def set_language(self, telegram_id: int, language: str) -> None: conn = self._require_conn() await conn.execute( """ UPDATE clients SET language = ?, updated_at = ? WHERE telegram_id = ? """, (language, _utcnow(), telegram_id), ) await conn.commit() async def save_message_link( self, client_telegram_id: int, client_message_id: int, manager_message_id: int, ) -> None: conn = self._require_conn() await conn.execute( "DELETE FROM message_links WHERE manager_message_id = ?", (manager_message_id,), ) await conn.execute( """ DELETE FROM message_links WHERE client_telegram_id = ? AND client_message_id = ? """, (client_telegram_id, client_message_id), ) await conn.execute( """ INSERT INTO message_links ( client_telegram_id, client_message_id, manager_message_id, created_at ) VALUES (?, ?, ?, ?) """, (client_telegram_id, client_message_id, manager_message_id, _utcnow()), ) await conn.commit() async def get_link_by_client_message( self, client_telegram_id: int, client_message_id: int, ) -> Optional[MessageLink]: conn = self._require_conn() async with conn.execute( """ SELECT client_telegram_id, client_message_id, manager_message_id FROM message_links WHERE client_telegram_id = ? AND client_message_id = ? """, (client_telegram_id, client_message_id), ) as cursor: row = await cursor.fetchone() if row is None: return None return MessageLink( client_telegram_id=row["client_telegram_id"], client_message_id=row["client_message_id"], manager_message_id=row["manager_message_id"], ) async def get_link_by_manager_message( self, manager_message_id: int, ) -> Optional[MessageLink]: conn = self._require_conn() async with conn.execute( """ SELECT client_telegram_id, client_message_id, manager_message_id FROM message_links WHERE manager_message_id = ? """, (manager_message_id,), ) as cursor: row = await cursor.fetchone() if row is None: return None return MessageLink( client_telegram_id=row["client_telegram_id"], client_message_id=row["client_message_id"], manager_message_id=row["manager_message_id"], ) async def delete_message_link( self, *, client_telegram_id: int, client_message_id: int, manager_message_id: int, ) -> None: conn = self._require_conn() await conn.execute( """ DELETE FROM message_links WHERE client_telegram_id = ? AND client_message_id = ? AND manager_message_id = ? """, (client_telegram_id, client_message_id, manager_message_id), ) await conn.commit()