"""Database layer — Postgres via asyncpg with a tiny chat-settings cache.""" from __future__ import annotations import logging import time from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, Optional import asyncpg from app.config import DEFAULTS LOGGER = logging.getLogger(__name__) SCHEMA_PATH = Path(__file__).parent.parent / "scripts" / "init_db.sql" @dataclass class ChatSettings: chat_id: int enabled: bool = DEFAULTS["enabled"] check_admins: bool = DEFAULTS["check_admins"] threshold: int = DEFAULTS["threshold"] filter_photos: bool = DEFAULTS["filter_photos"] filter_stickers: bool = DEFAULTS["filter_stickers"] filter_videos: bool = DEFAULTS["filter_videos"] filter_gifs: bool = DEFAULTS["filter_gifs"] filter_video_notes: bool = DEFAULTS["filter_video_notes"] filter_web_preview: bool = DEFAULTS["filter_web_preview"] block_porn: bool = DEFAULTS["block_porn"] block_hentai: bool = DEFAULTS["block_hentai"] block_sexy: bool = DEFAULTS["block_sexy"] antispam_enabled: bool = DEFAULTS["antispam_enabled"] warn_user: bool = DEFAULTS["warn_user"] max_violations: int = DEFAULTS["max_violations"] violation_window_minutes: int = DEFAULTS["violation_window_minutes"] ban_duration_minutes: int = DEFAULTS["ban_duration_minutes"] warn_delete_seconds: int = DEFAULTS["warn_delete_seconds"] delete_message_template: str = DEFAULTS["delete_message_template"] def blocked_categories(self) -> list[str]: out: list[str] = [] if self.block_porn: out.append("porn") if self.block_hentai: out.append("hentai") if self.block_sexy: out.append("sexy") return out @dataclass class _CacheEntry: settings: ChatSettings fetched_at: float = field(default_factory=time.monotonic) class Database: """Thin asyncpg wrapper. Pool is lazy.""" CACHE_TTL_SECONDS = 60.0 def __init__(self, dsn: str): self._dsn = dsn self._pool: Optional[asyncpg.Pool] = None self._cache: Dict[int, _CacheEntry] = {} async def connect(self) -> None: if self._pool is not None: return # Conservative pool — Render free has tight memory self._pool = await asyncpg.create_pool( dsn=self._dsn, min_size=1, max_size=4, command_timeout=10.0, ) await self._init_schema() LOGGER.info("Database pool ready") async def close(self) -> None: if self._pool is not None: await self._pool.close() self._pool = None async def _init_schema(self) -> None: sql = SCHEMA_PATH.read_text(encoding="utf-8") assert self._pool is not None async with self._pool.acquire() as conn: await conn.execute(sql) # ------------------------------------------------------------------ # Chat lifecycle # ------------------------------------------------------------------ async def upsert_chat( self, chat_id: int, title: Optional[str], added_by: Optional[int], member_count: Optional[int] = None, initial_status: str = "pending", ) -> str: """Insert chat or update title. Returns the resulting status string. Behaviour: - First insert uses ``initial_status`` (so the caller can pass 'approved' when the bot is added by the owner themselves). - On conflict (re-add) we keep prior status if it was 'approved' or 'rejected' (owner already decided), otherwise reset to 'pending'. """ assert self._pool is not None async with self._pool.acquire() as conn: row = await conn.fetchrow( """ INSERT INTO chats (chat_id, title, added_by, status, status_at, member_count) VALUES ($1, $2, $3, $4, NOW(), $5) ON CONFLICT (chat_id) DO UPDATE SET title = EXCLUDED.title, member_count = COALESCE(EXCLUDED.member_count, chats.member_count), updated_at = NOW(), status = CASE WHEN chats.status IN ('approved','rejected') THEN chats.status ELSE 'pending' END, status_at = CASE WHEN chats.status IN ('approved','rejected') THEN chats.status_at ELSE NOW() END RETURNING status """, chat_id, title, added_by, initial_status, member_count, ) await conn.execute( """ INSERT INTO chat_settings (chat_id) VALUES ($1) ON CONFLICT (chat_id) DO NOTHING """, chat_id, ) self._cache.pop(chat_id, None) return row["status"] if row else initial_status async def set_chat_status( self, chat_id: int, status: str, decided_by: Optional[int] = None, ) -> None: assert self._pool is not None if status not in ("pending", "approved", "rejected", "removed"): raise ValueError(f"Invalid chat status: {status}") async with self._pool.acquire() as conn: await conn.execute( """ UPDATE chats SET status = $1, status_at = NOW(), status_by = $2, updated_at = NOW() WHERE chat_id = $3 """, status, decided_by, chat_id, ) self._cache.pop(chat_id, None) async def get_chat(self, chat_id: int) -> Optional[Dict[str, Any]]: assert self._pool is not None async with self._pool.acquire() as conn: row = await conn.fetchrow( "SELECT chat_id, title, added_by, status, status_at, status_by, member_count " "FROM chats WHERE chat_id = $1", chat_id, ) return dict(row) if row else None async def is_chat_approved(self, chat_id: int) -> bool: chat = await self.get_chat(chat_id) return bool(chat and chat["status"] == "approved") async def list_chats_for_admin(self, admin_user_id: int) -> list[Dict[str, Any]]: """Best-effort list of chats that the bot knows about, used by /api/chats. We can't query Telegram membership cheaply for every chat, so we return chats where: - status is approved (so the user can pick from active ones), AND - the user is the original `added_by` OR the global owner. Membership re-verification happens per-chat on the dedicated endpoints. """ assert self._pool is not None async with self._pool.acquire() as conn: rows = await conn.fetch( """ SELECT chat_id, title, status, member_count, added_at, status_at FROM chats WHERE status = 'approved' ORDER BY status_at DESC """, ) return [dict(r) for r in rows] async def mark_chat_removed(self, chat_id: int) -> None: await self.set_chat_status(chat_id, "removed") async def delete_chat(self, chat_id: int) -> None: """Hard delete; kept for tests and admin tooling.""" assert self._pool is not None async with self._pool.acquire() as conn: await conn.execute("DELETE FROM chats WHERE chat_id = $1", chat_id) self._cache.pop(chat_id, None) # ------------------------------------------------------------------ # Settings # ------------------------------------------------------------------ async def get_settings(self, chat_id: int) -> ChatSettings: cached = self._cache.get(chat_id) if cached and (time.monotonic() - cached.fetched_at) < self.CACHE_TTL_SECONDS: return cached.settings assert self._pool is not None async with self._pool.acquire() as conn: row = await conn.fetchrow( "SELECT * FROM chat_settings WHERE chat_id = $1", chat_id, ) if row is None: settings = ChatSettings(chat_id=chat_id) else: settings = ChatSettings( chat_id=chat_id, enabled=row["enabled"], check_admins=row["check_admins"], threshold=row["threshold"], filter_photos=row["filter_photos"], filter_stickers=row["filter_stickers"], filter_videos=row["filter_videos"], filter_gifs=row["filter_gifs"], filter_video_notes=row["filter_video_notes"], filter_web_preview=row["filter_web_preview"], block_porn=row["block_porn"], block_hentai=row["block_hentai"], block_sexy=row["block_sexy"], antispam_enabled=row["antispam_enabled"], warn_user=row["warn_user"], max_violations=row["max_violations"], violation_window_minutes=row["violation_window_minutes"], ban_duration_minutes=row["ban_duration_minutes"], warn_delete_seconds=row["warn_delete_seconds"], delete_message_template=row["delete_message_template"], ) self._cache[chat_id] = _CacheEntry(settings=settings) return settings async def update_setting(self, chat_id: int, key: str, value: Any) -> None: # Whitelist of mutable columns to avoid SQL injection via dynamic key allowed = { "enabled", "check_admins", "threshold", "filter_photos", "filter_stickers", "filter_videos", "filter_gifs", "filter_video_notes", "filter_web_preview", "block_porn", "block_hentai", "block_sexy", "antispam_enabled", "warn_user", "max_violations", "violation_window_minutes", "ban_duration_minutes", "warn_delete_seconds", "delete_message_template", } if key not in allowed: raise ValueError(f"Cannot update unknown setting: {key}") assert self._pool is not None sql = ( f"UPDATE chat_settings SET {key} = $1, updated_at = NOW() " f"WHERE chat_id = $2" ) async with self._pool.acquire() as conn: await conn.execute(sql, value, chat_id) self._cache.pop(chat_id, None) # ------------------------------------------------------------------ # Violations # ------------------------------------------------------------------ async def record_violation( self, chat_id: int, user_id: int, category: Optional[str], score: Optional[float], media_kind: str, ) -> int: """Inserts a violation, returns count of recent violations within window.""" assert self._pool is not None async with self._pool.acquire() as conn: await conn.execute( """ INSERT INTO violations (chat_id, user_id, category, score, media_kind) VALUES ($1, $2, $3, $4, $5) """, chat_id, user_id, category, score, media_kind, ) settings = await self.get_settings(chat_id) window_minutes = settings.violation_window_minutes count = await conn.fetchval( """ SELECT COUNT(*) FROM violations WHERE chat_id = $1 AND user_id = $2 AND occurred_at > NOW() - ($3 || ' minutes')::INTERVAL """, chat_id, user_id, str(window_minutes), ) return int(count) # ------------------------------------------------------------------ # Events (timeline backing the mini-app stats + log screens) # ------------------------------------------------------------------ async def record_event( self, chat_id: int, kind: str, user_id: Optional[int] = None, user_name: Optional[str] = None, category: Optional[str] = None, score: Optional[float] = None, media_kind: Optional[str] = None, meta: Optional[Dict[str, Any]] = None, ) -> None: import json assert self._pool is not None async with self._pool.acquire() as conn: await conn.execute( """ INSERT INTO events (chat_id, kind, user_id, user_name, category, score, media_kind, meta) VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb) """, chat_id, kind, user_id, user_name, category, score, media_kind, json.dumps(meta) if meta else None, ) async def stats_for_chat(self, chat_id: int) -> Dict[str, Any]: """Aggregate stats for a single chat in one DB round-trip. Returns the shape consumed by /api/stats: d24, week, week_delta_pct, total, by_category, daily (last 7d), top_violators (last 30d top 5). """ assert self._pool is not None async with self._pool.acquire() as conn: row = await conn.fetchrow( """ WITH del AS ( SELECT * FROM events WHERE chat_id = $1 AND kind = 'delete' ) SELECT (SELECT COUNT(*) FROM del WHERE occurred_at > NOW() - INTERVAL '24 hours') AS d24, (SELECT COUNT(*) FROM del WHERE occurred_at > NOW() - INTERVAL '7 days') AS week, (SELECT COUNT(*) FROM del WHERE occurred_at > NOW() - INTERVAL '14 days' AND occurred_at <= NOW() - INTERVAL '7 days') AS prev_week, (SELECT COUNT(*) FROM del) AS total """, chat_id, ) cats = await conn.fetch( """ SELECT COALESCE(category,'unknown') AS key, COUNT(*) AS count FROM events WHERE chat_id = $1 AND kind = 'delete' AND occurred_at > NOW() - INTERVAL '7 days' GROUP BY 1 ORDER BY count DESC """, chat_id, ) daily = await conn.fetch( """ SELECT to_char(date_trunc('day', occurred_at), 'YYYY-MM-DD') AS day, COUNT(*) AS count FROM events WHERE chat_id = $1 AND kind = 'delete' AND occurred_at > NOW() - INTERVAL '7 days' GROUP BY 1 ORDER BY 1 """, chat_id, ) top = await conn.fetch( """ SELECT user_id, COALESCE(MAX(user_name), 'unknown') AS name, COUNT(*) AS count FROM events WHERE chat_id = $1 AND kind = 'delete' AND user_id IS NOT NULL AND occurred_at > NOW() - INTERVAL '30 days' GROUP BY user_id ORDER BY count DESC LIMIT 5 """, chat_id, ) d24 = int(row["d24"] or 0) week = int(row["week"] or 0) prev = int(row["prev_week"] or 0) total = int(row["total"] or 0) if prev > 0: delta = round((week - prev) * 100.0 / prev) elif week > 0: delta = 100 else: delta = 0 return { "d24": d24, "week": week, "week_delta_pct": delta, "total": total, "by_category": [{"key": r["key"], "count": int(r["count"])} for r in cats], "daily": [{"day": r["day"], "count": int(r["count"])} for r in daily], "top_violators": [ {"user_id": int(r["user_id"]), "name": r["name"], "count": int(r["count"])} for r in top ], } async def events_for_chat( self, chat_id: int, kinds: Optional[list[str]] = None, limit: int = 100, ) -> list[Dict[str, Any]]: assert self._pool is not None if kinds: sql = ( "SELECT id, occurred_at, kind, user_id, user_name, category, score, media_kind, meta " "FROM events WHERE chat_id = $1 AND kind = ANY($2::text[]) " "ORDER BY occurred_at DESC LIMIT $3" ) params: tuple = (chat_id, kinds, limit) else: sql = ( "SELECT id, occurred_at, kind, user_id, user_name, category, score, media_kind, meta " "FROM events WHERE chat_id = $1 ORDER BY occurred_at DESC LIMIT $2" ) params = (chat_id, limit) async with self._pool.acquire() as conn: rows = await conn.fetch(sql, *params) out = [] for r in rows: d = dict(r) if d.get("occurred_at") is not None: d["occurred_at"] = d["occurred_at"].isoformat() out.append(d) return out # ------------------------------------------------------------------ # User language (RU/AZ/EN) # ------------------------------------------------------------------ async def get_user_lang(self, user_id: int) -> str: assert self._pool is not None async with self._pool.acquire() as conn: row = await conn.fetchval( "SELECT lang FROM users WHERE user_id = $1", user_id, ) return row or "en" async def set_user_lang(self, user_id: int, lang: str) -> None: if lang not in ("ru", "az", "en"): raise ValueError(f"Unsupported language: {lang}") assert self._pool is not None async with self._pool.acquire() as conn: await conn.execute( """ INSERT INTO users (user_id, lang) VALUES ($1, $2) ON CONFLICT (user_id) DO UPDATE SET lang = EXCLUDED.lang, updated_at = NOW() """, user_id, lang, )