| """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) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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") |
|
|
