mirror of
https://github.com/artemium428/428th-exchange-bot.git
synced 2026-09-15 16:56:20 +00:00
281 lines
8.7 KiB
Python
281 lines
8.7 KiB
Python
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()
|