"""PostgreSQL data layer. Same tables as the Node version, so existing files and users carry over.""" from __future__ import annotations import ssl as _ssl from datetime import datetime from typing import Any from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import asyncpg SCHEMA = """ CREATE TABLE IF NOT EXISTS bot_files ( code TEXT NOT NULL, title TEXT NOT NULL, url TEXT NOT NULL, type TEXT NOT NULL DEFAULT 'document', sent INTEGER NOT NULL DEFAULT 0, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), delete_after_send BOOLEAN NOT NULL DEFAULT FALSE ); ALTER TABLE bot_files ADD COLUMN IF NOT EXISTS delete_after_send BOOLEAN NOT NULL DEFAULT FALSE; ALTER TABLE bot_files ADD COLUMN IF NOT EXISTS listed BOOLEAN NOT NULL DEFAULT TRUE; CREATE UNIQUE INDEX IF NOT EXISTS bot_files_code_lower ON bot_files (lower(code)); CREATE TABLE IF NOT EXISTS bot_users ( phone TEXT PRIMARY KEY, lang TEXT NOT NULL DEFAULT 'en', role TEXT NOT NULL DEFAULT 'free' CHECK (role IN ('free', 'premium', 'admin')), blocked BOOLEAN NOT NULL DEFAULT FALSE, daily_limit INTEGER, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), last_seen TIMESTAMPTZ NOT NULL DEFAULT now() ); CREATE TABLE IF NOT EXISTS bot_downloads ( id BIGSERIAL PRIMARY KEY, phone TEXT NOT NULL, code TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); CREATE INDEX IF NOT EXISTS bot_downloads_phone_time ON bot_downloads (phone, created_at); CREATE TABLE IF NOT EXISTS bot_sent_messages ( id BIGSERIAL PRIMARY KEY, jid TEXT NOT NULL, msg_id TEXT NOT NULL, delete_at TIMESTAMPTZ NOT NULL, attempts INTEGER NOT NULL DEFAULT 0, last_error TEXT ); CREATE INDEX IF NOT EXISTS bot_sent_messages_due ON bot_sent_messages (delete_at); CREATE TABLE IF NOT EXISTS bot_settings ( key TEXT PRIMARY KEY, value TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS bot_wa_session ( id TEXT PRIMARY KEY, data TEXT NOT NULL, updated_at TIMESTAMPTZ NOT NULL DEFAULT now() ); """ LOCAL_HOSTS = {"localhost", "127.0.0.1", "::1"} USER_FIELDS = {"role": "role", "blocked": "blocked", "dailyLimit": "daily_limit", "lang": "lang"} FILE_COLUMNS = ("code", "title", "url", "type", "sent", "listed", "delete_after_send", "created_at") def connection_settings(database_url: str, ssl_mode: str) -> tuple[str, Any]: """Strips sslmode= from the URL and returns (dsn, ssl argument for asyncpg). ssl_mode: '' (default: on for remote hosts, off for localhost and *.railway.internal), 'true' (verify certificate), 'insecure' (encrypt, skip certificate check), 'false' (off). """ parts = urlsplit(database_url) query = [ (k, v) for k, v in parse_qsl(parts.query, keep_blank_values=True) if k not in ("sslmode", "ssl", "uselibpqcompat", "sslrootcert") ] dsn = urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) host = parts.hostname or "" is_local = host in LOCAL_HOSTS or host.endswith(".railway.internal") mode = (ssl_mode or "").strip().lower() if mode in ("false", "0", "off"): ssl_arg = None elif mode == "insecure": ctx = _ssl.create_default_context() ctx.check_hostname = False ctx.verify_mode = _ssl.CERT_NONE ssl_arg = ctx elif mode in ("true", "1", "on") or (mode in ("", "auto") and not is_local): ssl_arg = _ssl.create_default_context() else: ssl_arg = None return dsn, ssl_arg def _iso(value: datetime | None) -> str | None: return value.isoformat() if value else None def _file(row: asyncpg.Record) -> dict: return { "code": row["code"], "title": row["title"], "url": row["url"], "type": row["type"], "sent": row["sent"], "listed": row["listed"], "deleteAfterSend": row["delete_after_send"], "createdAt": _iso(row["created_at"]), } def _user(row: asyncpg.Record) -> dict: user = { "id": row["phone"], "lang": row["lang"], "role": row["role"], "blocked": row["blocked"], "dailyLimit": row["daily_limit"], "createdAt": _iso(row["created_at"]), "lastSeen": _iso(row["last_seen"]), } if "today" in row.keys(): user["today"] = row["today"] return user class Database: def __init__(self, pool: asyncpg.Pool, default_auto_delete_hours: int = 24): self.pool = pool self.default_auto_delete_hours = default_auto_delete_hours @classmethod async def connect(cls, database_url: str, ssl_mode: str, default_auto_delete_hours: int = 24) -> "Database": dsn, ssl_arg = connection_settings(database_url, ssl_mode) pool = await asyncpg.create_pool(dsn, ssl=ssl_arg, min_size=1, max_size=5) return cls(pool, default_auto_delete_hours) async def close(self) -> None: await self.pool.close() async def init(self) -> None: async with self.pool.acquire() as conn: await conn.execute(SCHEMA) # ----- files --------------------------------------------------------- async def list_files(self) -> list[dict]: rows = await self.pool.fetch("SELECT * FROM bot_files ORDER BY created_at, code") return [_file(r) for r in rows] async def find_file(self, code: str) -> dict | None: row = await self.pool.fetchrow("SELECT * FROM bot_files WHERE lower(code) = lower($1)", code) return _file(row) if row else None async def add_file(self, code: str, title: str, url: str, type_: str, delete_after_send: bool = False) -> dict: row = await self.pool.fetchrow( "INSERT INTO bot_files (code, title, url, type, delete_after_send) VALUES ($1, $2, $3, $4, $5) RETURNING *", code, title, url, type_, bool(delete_after_send), ) return _file(row) async def delete_file(self, code: str) -> int: status = await self.pool.execute("DELETE FROM bot_files WHERE lower(code) = lower($1)", code) return int(status.split()[-1]) async def set_listed(self, code: str, listed: bool) -> dict | None: row = await self.pool.fetchrow( "UPDATE bot_files SET listed = $2 WHERE lower(code) = lower($1) RETURNING *", code, bool(listed) ) return _file(row) if row else None async def list_catalog(self) -> list[dict]: rows = await self.pool.fetch("SELECT code, title, type FROM bot_files WHERE listed ORDER BY lower(title)") return [{"code": r["code"], "title": r["title"], "type": r["type"]} for r in rows] async def record_send(self, code: str) -> bool: """Counts a send. For one-time files, deletes the file in the same step. Returns True if deleted.""" await self.pool.execute("UPDATE bot_files SET sent = sent + 1 WHERE lower(code) = lower($1)", code) status = await self.pool.execute( "DELETE FROM bot_files WHERE lower(code) = lower($1) AND delete_after_send = TRUE", code ) return int(status.split()[-1]) > 0 # ----- users --------------------------------------------------------- async def touch_user(self, phone: str, detected_lang: str | None = None) -> tuple[dict, bool]: row = await self.pool.fetchrow( """INSERT INTO bot_users (phone, lang) VALUES ($1::text, COALESCE($2::text, 'en')) ON CONFLICT (phone) DO UPDATE SET last_seen = now(), lang = COALESCE($2::text, bot_users.lang) RETURNING *, (xmax = 0) AS inserted""", phone, detected_lang, ) return _user(row), bool(row["inserted"]) async def get_user(self, phone: str) -> dict | None: row = await self.pool.fetchrow("SELECT * FROM bot_users WHERE phone = $1", phone) return _user(row) if row else None async def list_users(self, q: str = "") -> list[dict]: rows = await self.pool.fetch( """SELECT u.*, (SELECT count(*)::int FROM bot_downloads d WHERE d.phone = u.phone AND d.created_at > now() - interval '24 hours') AS today FROM bot_users u WHERE $1::text = '' OR u.phone LIKE '%' || $1::text || '%' ORDER BY u.last_seen DESC LIMIT 200""", q, ) return [_user(r) for r in rows] async def update_user(self, phone: str, patch: dict) -> dict | None: sets, vals = [], [] for key, col in USER_FIELDS.items(): if key in patch: vals.append(patch[key]) sets.append(f"{col} = ${len(vals)}") if not sets: return None vals.append(phone) row = await self.pool.fetchrow( f"UPDATE bot_users SET {', '.join(sets)} WHERE phone = ${len(vals)} RETURNING *", *vals ) return _user(row) if row else None async def downloads_today(self, phone: str) -> int: return await self.pool.fetchval( "SELECT count(*)::int FROM bot_downloads WHERE phone = $1 AND created_at > now() - interval '24 hours'", phone, ) async def record_download(self, phone: str, code: str) -> None: await self.pool.execute("INSERT INTO bot_downloads (phone, code) VALUES ($1, $2)", phone, code) async def stats(self) -> dict: row = await self.pool.fetchrow( """SELECT (SELECT count(*)::int FROM bot_users) AS users, (SELECT count(*)::int FROM bot_users WHERE role = 'premium') AS premium, (SELECT count(*)::int FROM bot_users WHERE role = 'admin') AS admins, (SELECT count(*)::int FROM bot_users WHERE blocked) AS blocked, (SELECT count(*)::int FROM bot_downloads WHERE created_at > now() - interval '24 hours') AS downloads_today, (SELECT count(*)::int FROM bot_downloads) AS downloads_total, (SELECT count(*)::int FROM bot_files) AS files""" ) return dict(row) # ----- settings ------------------------------------------------------ def default_settings(self) -> dict: return {"limits": {"free": 3, "premium": 30, "admin": None}, "autoDeleteHours": self.default_auto_delete_hours} async def get_settings(self) -> dict: import json defaults = self.default_settings() raw_text = await self.pool.fetchval("SELECT value FROM bot_settings WHERE key = 'main'") raw = json.loads(raw_text) if raw_text else {} return { "limits": {**defaults["limits"], **(raw.get("limits") or {})}, "autoDeleteHours": raw.get("autoDeleteHours", defaults["autoDeleteHours"]), } async def set_settings(self, patch: dict) -> dict: import json current = await self.get_settings() nxt = { "limits": {**current["limits"], **(patch.get("limits") or {})}, "autoDeleteHours": patch.get("autoDeleteHours", current["autoDeleteHours"]), } await self.pool.execute( """INSERT INTO bot_settings (key, value) VALUES ('main', $1) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value""", json.dumps(nxt), ) return nxt # ----- scheduled deletes ---------------------------------------------- async def schedule_deletes(self, items: list[tuple[str, str]], when: datetime) -> None: for jid, msg_id in items: if not msg_id: continue await self.pool.execute( "INSERT INTO bot_sent_messages (jid, msg_id, delete_at) VALUES ($1, $2, $3)", jid, msg_id, when ) async def due_deletes(self, limit: int = 20) -> list[asyncpg.Record]: return await self.pool.fetch( "SELECT id, jid, msg_id FROM bot_sent_messages WHERE delete_at <= now() ORDER BY delete_at LIMIT $1", limit, ) async def remove_delete(self, row_id: int) -> None: await self.pool.execute("DELETE FROM bot_sent_messages WHERE id = $1", row_id) async def fail_delete(self, row_id: int, message: str) -> None: """Retry in 5 minutes. Gives up after 5 failed attempts.""" await self.pool.execute( """UPDATE bot_sent_messages SET attempts = attempts + 1, last_error = $2, delete_at = now() + interval '5 minutes' WHERE id = $1""", row_id, str(message)[:300], ) await self.pool.execute("DELETE FROM bot_sent_messages WHERE id = $1 AND attempts >= 5", row_id) # ----- WhatsApp session (encrypted blob, see crypto.py) --------------- async def save_session(self, sealed: str) -> None: await self.pool.execute( """INSERT INTO bot_wa_session (id, data, updated_at) VALUES ('main', $1, now()) ON CONFLICT (id) DO UPDATE SET data = EXCLUDED.data, updated_at = now()""", sealed, ) async def load_session(self) -> str | None: return await self.pool.fetchval("SELECT data FROM bot_wa_session WHERE id = 'main'") async def clear_session(self) -> None: await self.pool.execute("DELETE FROM bot_wa_session")