137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
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)
|