import json import asyncpg from pyrogram.types import User from utils.read.models import AccountView ACCOUNTS_CHANGED_CHANNEL = "accounts_changed" _ACCOUNT_COLS = "account_id, label, phone, tg_user_id, is_active, device_model" _UPSERT_ACCOUNT = """ INSERT INTO accounts (tg_user_id, label, phone, session_name, is_active, raw, updated_at) VALUES ($1, $2, $3, $4, TRUE, $5::jsonb, now()) ON CONFLICT (tg_user_id) DO UPDATE SET label = EXCLUDED.label, phone = EXCLUDED.phone, session_name = EXCLUDED.session_name, is_active = TRUE, raw = EXCLUDED.raw, updated_at = now() RETURNING account_id """ _SET_DEVICE_MODEL = f""" UPDATE accounts SET device_model = $2, updated_at = now() WHERE account_id = $1 RETURNING {_ACCOUNT_COLS} """ # noqa: S608 _self_ids: dict[int, int] = {} async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None: cached = _self_ids.get(account_id) if cached is not None: return cached tg_user_id = await pool.fetchval( "SELECT tg_user_id FROM accounts WHERE account_id = $1", account_id ) if tg_user_id is not None: _self_ids[account_id] = tg_user_id return tg_user_id _known_ids: set[int] = set() def _account_label(row: asyncpg.Record) -> str: named = f", {row['label']}" if row["label"] else "" return f"account_id={row['account_id']} (Telegram id {row['tg_user_id']}{named})" async def unknown_account_hint(pool: asyncpg.Pool, account_id: int) -> str | None: """Explain an account_id that matches no account, or None if it is valid.""" if account_id in _known_ids: return None rows = await pool.fetch( "SELECT account_id, tg_user_id, label FROM accounts ORDER BY account_id" ) _known_ids.update(row["account_id"] for row in rows) if account_id in _known_ids: return None known = "; ".join(_account_label(row) for row in rows) or "none" mistaken = next((row for row in rows if row["tg_user_id"] == account_id), None) if mistaken is not None: return ( f"{account_id} is a Telegram user id, not an account_id. In beavergram " f"that account is account_id={mistaken['account_id']}. " f"Retry with account_id={mistaken['account_id']}." ) return f"No account with account_id={account_id}. Archived accounts: {known}." async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]: rows = await pool.fetch( f"SELECT {_ACCOUNT_COLS} FROM accounts ORDER BY account_id" # noqa: S608 ) return [AccountView(**dict(row)) for row in rows] async def get_account(pool: asyncpg.Pool, account_id: int) -> AccountView | None: row = await pool.fetchrow( f"SELECT {_ACCOUNT_COLS} FROM accounts WHERE account_id = $1", # noqa: S608 account_id, ) return AccountView(**dict(row)) if row else None async def sync_account(pool: asyncpg.Pool, me: User, session_name: str) -> int: raw = json.dumps( { "id": me.id, "first_name": me.first_name, "last_name": me.last_name, "username": me.username, "phone_number": me.phone_number, } ) label = " ".join(filter(None, [me.first_name, me.last_name])) or me.username return await pool.fetchval( _UPSERT_ACCOUNT, me.id, label, me.phone_number, session_name, raw ) async def set_device_model( pool: asyncpg.Pool, account_id: int, device_model: str ) -> AccountView | None: row = await pool.fetchrow(_SET_DEVICE_MODEL, account_id, device_model) return AccountView(**dict(row)) if row else None async def session_device_models(pool: asyncpg.Pool) -> dict[str, str]: rows = await pool.fetch( "SELECT session_name, device_model FROM accounts WHERE device_model IS NOT NULL" ) return {row["session_name"]: row["device_model"] for row in rows} async def deactivate_account(pool: asyncpg.Pool, account_id: int) -> str | None: return await pool.fetchval( "UPDATE accounts SET is_active = FALSE, updated_at = now() " "WHERE account_id = $1 RETURNING session_name", account_id, ) async def inactive_session_names(pool: asyncpg.Pool) -> set[str]: rows = await pool.fetch("SELECT session_name FROM accounts WHERE NOT is_active") return {row["session_name"] for row in rows} async def notify_accounts_changed(pool: asyncpg.Pool) -> None: await pool.execute("SELECT pg_notify($1, '')", ACCOUNTS_CHANGED_CHANNEL)