perf(*): bound list queries, batch fetches, harden media downloads
This commit is contained in:
@@ -0,0 +1,49 @@
|
||||
"""read path indexes
|
||||
|
||||
Revision ID: a6d2f8b41c9e
|
||||
Revises: f2b8d3c9a51e
|
||||
Create Date: 2026-09-01 12:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
revision: str = "a6d2f8b41c9e"
|
||||
down_revision: str | None = "f2b8d3c9a51e"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("CREATE INDEX ix_jobs_account_id ON jobs (account_id, id DESC)")
|
||||
op.execute(
|
||||
"CREATE INDEX ix_messages_deleted ON messages "
|
||||
"(account_id, deleted_at DESC) WHERE deleted_at IS NOT NULL"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_messages_pinned ON messages "
|
||||
"(account_id, chat_id, date DESC, message_id DESC) "
|
||||
"WHERE raw->>'service' LIKE '%PINNED_MESSAGE%'"
|
||||
)
|
||||
op.execute("CREATE INDEX ix_watches_account ON watches (account_id, id DESC)")
|
||||
op.execute("CREATE INDEX ix_alerts_account_ts ON alerts (account_id, ts DESC)")
|
||||
op.execute(
|
||||
"CREATE INDEX ix_stories_account_date ON stories "
|
||||
"(account_id, date DESC NULLS LAST)"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX ix_annotations_account_msg ON annotations "
|
||||
"(account_id, chat_id, message_id)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP INDEX IF EXISTS ix_annotations_account_msg")
|
||||
op.execute("DROP INDEX IF EXISTS ix_stories_account_date")
|
||||
op.execute("DROP INDEX IF EXISTS ix_alerts_account_ts")
|
||||
op.execute("DROP INDEX IF EXISTS ix_watches_account")
|
||||
op.execute("DROP INDEX IF EXISTS ix_messages_pinned")
|
||||
op.execute("DROP INDEX IF EXISTS ix_messages_deleted")
|
||||
op.execute("DROP INDEX IF EXISTS ix_jobs_account_id")
|
||||
@@ -8,6 +8,7 @@ from fastapi import APIRouter, HTTPException, Query
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils.jobs import enqueue
|
||||
from utils.read.models import DEFAULT_LIMIT, Page
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["backfill"], route_class=DishkaRoute)
|
||||
|
||||
@@ -130,16 +131,26 @@ async def list_jobs(
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
account_id: Annotated[int, Query()],
|
||||
status: Annotated[str | None, Query()] = None,
|
||||
limit: Annotated[int, Query()] = DEFAULT_LIMIT,
|
||||
offset: Annotated[int, Query()] = 0,
|
||||
) -> list[JobView]:
|
||||
page = Page(limit=limit, offset=offset)
|
||||
if status is None:
|
||||
rows = await pool.fetch(
|
||||
"SELECT * FROM jobs WHERE account_id = $1 ORDER BY id DESC", account_id
|
||||
"SELECT * FROM jobs WHERE account_id = $1 "
|
||||
"ORDER BY id DESC LIMIT $2 OFFSET $3",
|
||||
account_id,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
)
|
||||
else:
|
||||
rows = await pool.fetch(
|
||||
"SELECT * FROM jobs WHERE account_id = $1 AND status = $2 ORDER BY id DESC",
|
||||
"SELECT * FROM jobs WHERE account_id = $1 AND status = $2 "
|
||||
"ORDER BY id DESC LIMIT $3 OFFSET $4",
|
||||
account_id,
|
||||
status,
|
||||
page.capped_limit,
|
||||
page.offset,
|
||||
)
|
||||
return [_to_view(row) for row in rows]
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@ from utils.policy import repository
|
||||
from utils.read import chats
|
||||
from utils.read.models import (
|
||||
DEFAULT_LIMIT,
|
||||
MAX_LIMIT,
|
||||
ChatListItem,
|
||||
MessageVersionView,
|
||||
MessageView,
|
||||
@@ -49,6 +50,14 @@ async def list_chats(
|
||||
)
|
||||
|
||||
|
||||
@router.get("/chats/batch")
|
||||
async def get_chats_batch(
|
||||
pool: FromDishka[asyncpg.Pool], account_id: AccountId, ids: Annotated[str, Query()]
|
||||
) -> list[ChatListItem]:
|
||||
parsed = [int(part) for part in ids.split(",") if part.strip()]
|
||||
return await chats.get_chats(pool, account_id, parsed[:MAX_LIMIT])
|
||||
|
||||
|
||||
@router.get("/chats/{chat_id}")
|
||||
async def get_chat(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
||||
|
||||
@@ -13,7 +13,7 @@ from utils.policy import repository as policy_repository
|
||||
from utils.policy.defaults import TRACKING
|
||||
from utils.policy.models import ScopeType
|
||||
from utils.read import discover
|
||||
from utils.read.models import DiscoverItem
|
||||
from utils.read.models import MAX_LIMIT, DiscoverItem
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["discover"], route_class=DishkaRoute)
|
||||
|
||||
@@ -68,6 +68,7 @@ async def discover_peers(
|
||||
) -> list[DiscoverItem]:
|
||||
if not query.strip():
|
||||
return []
|
||||
limit = min(limit, MAX_LIMIT)
|
||||
items = await discover.search(pool, account_id, query, limit)
|
||||
if not remote:
|
||||
return items
|
||||
|
||||
@@ -52,9 +52,15 @@ async def chat_links(
|
||||
|
||||
@router.get("/calendar")
|
||||
async def chat_calendar(
|
||||
pool: FromDishka[asyncpg.Pool], chat_id: int, account_id: AccountId
|
||||
pool: FromDishka[asyncpg.Pool],
|
||||
chat_id: int,
|
||||
account_id: AccountId,
|
||||
date_from: Annotated[datetime | None, Query()] = None,
|
||||
date_to: Annotated[datetime | None, Query()] = None,
|
||||
) -> list[DayCount]:
|
||||
return await profile.daily_counts(pool, account_id, chat_id)
|
||||
return await profile.daily_counts(
|
||||
pool, account_id, chat_id, date_from=date_from, date_to=date_to
|
||||
)
|
||||
|
||||
|
||||
@router.get("/message-at")
|
||||
|
||||
@@ -122,7 +122,7 @@ async def max_message_id(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
) -> int | None:
|
||||
return await pool.fetchval(
|
||||
"SELECT max(message_id) FROM messages WHERE account_id = $1 AND chat_id = $2",
|
||||
"SELECT last_message_id FROM chat_stats WHERE account_id = $1 AND chat_id = $2",
|
||||
account_id,
|
||||
chat_id,
|
||||
)
|
||||
|
||||
@@ -31,6 +31,7 @@ class PyroClient(Client):
|
||||
lang_code="en",
|
||||
system_lang_code="en-US",
|
||||
client_platform=enums.ClientPlatform.DESKTOP,
|
||||
max_concurrent_transmissions=4,
|
||||
)
|
||||
self.capture: CaptureContext | None = None
|
||||
|
||||
|
||||
@@ -50,7 +50,7 @@ async def _reset_media_sessions(client: Client) -> None:
|
||||
async def download_bytes(client: Client, target: Downloadable) -> bytes | None:
|
||||
try:
|
||||
buffer = await client.download_media(target, in_memory=True)
|
||||
except TimeoutError:
|
||||
except (TimeoutError, OSError):
|
||||
await _reset_media_sessions(client)
|
||||
buffer = await client.download_media(target, in_memory=True)
|
||||
return buffer.getvalue() if isinstance(buffer, BytesIO) else None
|
||||
|
||||
@@ -69,6 +69,9 @@ class JobConsumer:
|
||||
requeued = await repository.requeue_running(self.pool, self.account_id)
|
||||
if requeued:
|
||||
logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
|
||||
pruned = await repository.prune_finished(self.pool, self.account_id)
|
||||
if pruned:
|
||||
logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]")
|
||||
conn = await asyncpg.connect(dsn=env.db.connection_url)
|
||||
await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set())
|
||||
logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import PeerIdInvalid
|
||||
from pyrogram.errors import FloodPremiumWait, FloodWait, PeerIdInvalid
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot.modules.capture import capture_message
|
||||
@@ -11,6 +11,7 @@ from userbot.modules.jobs.registry import register
|
||||
from userbot.modules.stt import repository as stt_repo
|
||||
from userbot.modules.stt import should_transcribe_on_backfill
|
||||
from userbot.modules.stt.gate import safe_transcribe
|
||||
from utils.logging import logger
|
||||
from utils.policy.models import CaptureToggles
|
||||
|
||||
SAVE_EVERY = 100
|
||||
@@ -61,21 +62,30 @@ async def backfill(ctx: JobContext) -> None:
|
||||
min_id = await resolve_min_id(ctx, chat_id)
|
||||
await ctx.save_cursor({"max_id": max_id, "min_id": min_id})
|
||||
processed = ctx.job.progress.get("processed", 0)
|
||||
errors = ctx.job.progress.get("errors", 0)
|
||||
self_id = client.me.id if client.me else None
|
||||
try:
|
||||
async for message in client.get_chat_history(
|
||||
chat_id, max_id=max_id, min_id=min_id
|
||||
):
|
||||
await capture_message(client, message, capture, toggles)
|
||||
await maybe_transcribe(client, capture, chat_id, message, self_id)
|
||||
try:
|
||||
await capture_message(client, message, capture, toggles)
|
||||
await maybe_transcribe(client, capture, chat_id, message, self_id)
|
||||
except (FloodWait, FloodPremiumWait, PeerIdInvalid):
|
||||
raise
|
||||
except Exception:
|
||||
errors += 1
|
||||
logger.exception(f"Backfill capture failed for {chat_id}/{message.id}")
|
||||
processed += 1
|
||||
if processed % SAVE_EVERY == 0:
|
||||
next_max = message.id - 1
|
||||
await ctx.save_cursor({"max_id": next_max, "min_id": min_id})
|
||||
await ctx.report_progress({"processed": processed, "max_id": next_max})
|
||||
await ctx.report_progress(
|
||||
{"processed": processed, "errors": errors, "max_id": next_max}
|
||||
)
|
||||
if await ctx.is_canceled():
|
||||
return
|
||||
except PeerIdInvalid:
|
||||
await ctx.report_progress({"processed": processed, "error": "peer_id_invalid"})
|
||||
return
|
||||
await ctx.report_progress({"processed": processed, "done": True})
|
||||
await ctx.report_progress({"processed": processed, "errors": errors, "done": True})
|
||||
|
||||
@@ -65,16 +65,28 @@ async def _enrich_members(client: Client, ctx: CaptureContext, chat_id: int) ->
|
||||
return
|
||||
|
||||
|
||||
async def _fetch_users(client: Client, ids: list[int]) -> list[User]:
|
||||
try:
|
||||
users = await client.get_users(ids)
|
||||
except BadRequest:
|
||||
users = []
|
||||
for sender_id in ids:
|
||||
try:
|
||||
users.append(await client.get_users(sender_id))
|
||||
except BadRequest:
|
||||
continue
|
||||
if isinstance(users, User):
|
||||
return [users]
|
||||
return [user for user in users if isinstance(user, User)]
|
||||
|
||||
|
||||
async def _enrich_senders(client: Client, ctx: CaptureContext, chat_id: int) -> None:
|
||||
rows = await ctx.pool.fetch(_MISSING_SENDERS, ctx.account_id, chat_id)
|
||||
ids = [row["sender_id"] for row in rows]
|
||||
for sender_id in ids:
|
||||
try:
|
||||
user = await client.get_users(sender_id)
|
||||
except BadRequest:
|
||||
continue
|
||||
if isinstance(user, User):
|
||||
await _save_user(ctx, user)
|
||||
if not ids:
|
||||
return
|
||||
for user in await _fetch_users(client, ids):
|
||||
await _save_user(ctx, user)
|
||||
|
||||
|
||||
@register("enrich_chat")
|
||||
|
||||
@@ -18,6 +18,10 @@ async def fetch_media(ctx: JobContext) -> None:
|
||||
if isinstance(message, list):
|
||||
message = message[0] if message else None
|
||||
if message is None or message.empty:
|
||||
return
|
||||
msg = f"message {chat_id}/{message_id} is unavailable"
|
||||
raise RuntimeError(msg)
|
||||
toggles = CaptureToggles(media=True, self_destruct_media=True)
|
||||
await capture_media(client, message, capture, chat_id, message_id, toggles)
|
||||
ok = await capture_media(client, message, capture, chat_id, message_id, toggles)
|
||||
if not ok:
|
||||
msg = f"media download failed for {chat_id}/{message_id}"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
@@ -48,6 +48,16 @@ async def requeue_running(pool: asyncpg.Pool, account_id: int) -> int:
|
||||
return int(result.split()[-1])
|
||||
|
||||
|
||||
async def prune_finished(pool: asyncpg.Pool, account_id: int) -> int:
|
||||
result = await pool.execute(
|
||||
"DELETE FROM jobs WHERE account_id = $1 "
|
||||
"AND status IN ('done', 'canceled', 'failed') "
|
||||
"AND finished_at < now() - interval '30 days'",
|
||||
account_id,
|
||||
)
|
||||
return int(result.split()[-1])
|
||||
|
||||
|
||||
async def save_cursor(pool: asyncpg.Pool, job_id: int, cursor: dict[str, Any]) -> None:
|
||||
await pool.execute(
|
||||
"UPDATE jobs SET cursor = $2::jsonb, updated_at = now() WHERE id = $1",
|
||||
|
||||
@@ -1,14 +1,23 @@
|
||||
from typing import Any
|
||||
|
||||
from pyrogram import Client
|
||||
from pyrogram.errors import (
|
||||
FileIdInvalid,
|
||||
FileReferenceExpired,
|
||||
FileReferenceInvalid,
|
||||
FloodPremiumWait,
|
||||
FloodWait,
|
||||
)
|
||||
from pyrogram.types import Message
|
||||
|
||||
from userbot.modules.capture import repository
|
||||
from userbot.modules.capture.context import CaptureContext
|
||||
from userbot.modules.download import download_bytes
|
||||
from userbot.modules.download import Downloadable, download_bytes
|
||||
from utils.logging import logger
|
||||
from utils.policy.models import CaptureToggles
|
||||
|
||||
STALE_FILE_ID = (FileIdInvalid, FileReferenceExpired, FileReferenceInvalid)
|
||||
|
||||
_MEDIA_ATTRS = (
|
||||
"photo",
|
||||
"video",
|
||||
@@ -48,6 +57,48 @@ def media_unique_id(message: Message) -> str | None:
|
||||
return getattr(obj, "file_unique_id", None) if obj is not None else None
|
||||
|
||||
|
||||
def _download_target(message: Message) -> Downloadable | None:
|
||||
kind, obj = media_object(message)
|
||||
if obj is None:
|
||||
return None
|
||||
return message if getattr(message, kind or "", None) is obj else obj
|
||||
|
||||
|
||||
async def _fresh_target(
|
||||
client: Client, chat_id: int, message_id: int
|
||||
) -> Downloadable | None:
|
||||
message = await client.get_messages(chat_id, message_id)
|
||||
if not isinstance(message, Message) or message.empty:
|
||||
return None
|
||||
return _download_target(message)
|
||||
|
||||
|
||||
async def _download(
|
||||
client: Client, target: Downloadable | None, chat_id: int, message_id: int
|
||||
) -> bytes | None:
|
||||
if target is None:
|
||||
return None
|
||||
try:
|
||||
return await download_bytes(client, target)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
|
||||
)
|
||||
return None
|
||||
except STALE_FILE_ID:
|
||||
fresh = await _fresh_target(client, chat_id, message_id)
|
||||
if fresh is None:
|
||||
return None
|
||||
try:
|
||||
return await download_bytes(client, fresh)
|
||||
except (TimeoutError, *STALE_FILE_ID):
|
||||
logger.warning(
|
||||
f"[yellow]Media download failed after refetch "
|
||||
f"for {chat_id}/{message_id}.[/]"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def capture_media( # noqa: PLR0913
|
||||
client: Client,
|
||||
message: Message,
|
||||
@@ -55,10 +106,10 @@ async def capture_media( # noqa: PLR0913
|
||||
chat_id: int,
|
||||
message_id: int,
|
||||
toggles: CaptureToggles,
|
||||
) -> None:
|
||||
) -> bool:
|
||||
kind, obj = media_object(message)
|
||||
if obj is None:
|
||||
return
|
||||
return True
|
||||
unique_id = getattr(obj, "file_unique_id", None)
|
||||
ttl = getattr(obj, "ttl_seconds", None)
|
||||
want = toggles.self_destruct_media if ttl else toggles.media
|
||||
@@ -66,6 +117,7 @@ async def capture_media( # noqa: PLR0913
|
||||
mime = getattr(obj, "mime_type", None)
|
||||
storage_key: str | None = None
|
||||
downloaded = False
|
||||
flood: FloodWait | FloodPremiumWait | None = None
|
||||
if want:
|
||||
existing = await repository.current_media(
|
||||
ctx.pool, ctx.account_id, chat_id, message_id
|
||||
@@ -80,13 +132,11 @@ async def capture_media( # noqa: PLR0913
|
||||
file_size = existing["file_size"]
|
||||
downloaded = True
|
||||
else:
|
||||
target = message if getattr(message, kind or "", None) is obj else obj
|
||||
target = _download_target(message)
|
||||
try:
|
||||
data = await download_bytes(client, target)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
f"[yellow]Media download timed out for {chat_id}/{message_id}.[/]"
|
||||
)
|
||||
data = await _download(client, target, chat_id, message_id)
|
||||
except (FloodWait, FloodPremiumWait) as exc:
|
||||
flood = exc
|
||||
data = None
|
||||
if data is not None:
|
||||
storage_key = ctx.storage.put(data)
|
||||
@@ -105,3 +155,6 @@ async def capture_media( # noqa: PLR0913
|
||||
unique_id,
|
||||
downloaded=downloaded,
|
||||
)
|
||||
if flood is not None:
|
||||
raise flood
|
||||
return downloaded or not want
|
||||
|
||||
@@ -49,18 +49,6 @@ async def _cancel(task: asyncio.Task) -> None:
|
||||
await task
|
||||
|
||||
|
||||
async def _enqueue_once(pool: asyncpg.Pool, account_id: int, kind: str) -> None:
|
||||
existing = await pool.fetchval(
|
||||
"SELECT 1 FROM jobs WHERE account_id = $1 AND kind = $2 "
|
||||
"AND status IN ('pending', 'running') LIMIT 1",
|
||||
account_id,
|
||||
kind,
|
||||
)
|
||||
if existing is None:
|
||||
await enqueue(pool, account_id, kind, {})
|
||||
logger.info(f"[green]Queued {kind}.[/]")
|
||||
|
||||
|
||||
class AccountRegistry:
|
||||
def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
|
||||
self._pool = pool
|
||||
@@ -121,8 +109,8 @@ class AccountRegistry:
|
||||
client, asyncio.create_task(consumer.run()), device_model
|
||||
)
|
||||
logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
|
||||
await _enqueue_once(self._pool, account_id, "sync_dialogs")
|
||||
await _enqueue_once(self._pool, account_id, "sync_contacts")
|
||||
await enqueue(self._pool, account_id, "sync_dialogs", {})
|
||||
await enqueue(self._pool, account_id, "sync_contacts", {})
|
||||
|
||||
async def _stop(self, session_name: str) -> None:
|
||||
account = self._running.pop(session_name, None)
|
||||
|
||||
@@ -9,12 +9,22 @@ JOBS_CHANGED_CHANNEL = "jobs_changed"
|
||||
async def enqueue(
|
||||
pool: asyncpg.Pool, account_id: int, kind: str, params: dict[str, Any]
|
||||
) -> int:
|
||||
params_json = json.dumps(params)
|
||||
existing = await pool.fetchval(
|
||||
"SELECT id FROM jobs WHERE account_id = $1 AND kind = $2 "
|
||||
"AND params = $3::jsonb AND status IN ('pending', 'running') LIMIT 1",
|
||||
account_id,
|
||||
kind,
|
||||
params_json,
|
||||
)
|
||||
if existing is not None:
|
||||
return existing
|
||||
job_id = await pool.fetchval(
|
||||
"INSERT INTO jobs (account_id, kind, params) "
|
||||
"VALUES ($1, $2, $3::jsonb) RETURNING id",
|
||||
account_id,
|
||||
kind,
|
||||
json.dumps(params),
|
||||
params_json,
|
||||
)
|
||||
await pool.execute(f"NOTIFY {JOBS_CHANGED_CHANNEL}")
|
||||
return job_id
|
||||
|
||||
@@ -30,10 +30,19 @@ RETURNING {_ACCOUNT_COLS}
|
||||
""" # noqa: S608
|
||||
|
||||
|
||||
_self_ids: dict[int, int] = {}
|
||||
|
||||
|
||||
async def self_user_id(pool: asyncpg.Pool, account_id: int) -> int | None:
|
||||
return await pool.fetchval(
|
||||
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
|
||||
|
||||
|
||||
async def list_accounts(pool: asyncpg.Pool) -> list[AccountView]:
|
||||
|
||||
@@ -33,7 +33,7 @@ async def message_volume(
|
||||
|
||||
|
||||
async def response_stats(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, *, days: int = 365
|
||||
) -> ResponseStats:
|
||||
self_id = await self_user_id(pool, account_id)
|
||||
rows = await pool.fetch(
|
||||
@@ -43,6 +43,7 @@ async def response_stats(
|
||||
"lag(date) OVER w AS prev_date "
|
||||
"FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND sender_id IS NOT NULL "
|
||||
"AND date >= $4 "
|
||||
"WINDOW w AS (ORDER BY date, message_id)), "
|
||||
"resp AS ("
|
||||
"SELECT (sender_id = $3) AS is_mine, "
|
||||
@@ -55,6 +56,7 @@ async def response_stats(
|
||||
account_id,
|
||||
chat_id,
|
||||
self_id,
|
||||
datetime.now(UTC) - timedelta(days=days),
|
||||
)
|
||||
stats = ResponseStats(
|
||||
mine_median_seconds=None, mine_count=0, their_median_seconds=None, their_count=0
|
||||
|
||||
@@ -21,7 +21,7 @@ WHERE account_id = $1 AND owner_id = $2 AND unique_id = $3
|
||||
_AVATAR_HISTORY = """
|
||||
SELECT unique_id, first_seen_at, downloaded FROM avatars
|
||||
WHERE account_id = $1 AND owner_id = $2
|
||||
ORDER BY first_seen_at DESC
|
||||
ORDER BY first_seen_at DESC LIMIT 200
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -62,6 +62,18 @@ _ONE_ID = """
|
||||
WHERE account_id = $1 AND scope_type = 'chat' AND scope_id = $2
|
||||
"""
|
||||
|
||||
_MANY_IDS = """
|
||||
SELECT chat_id FROM chat_stats
|
||||
WHERE account_id = $1 AND chat_id = ANY($2::bigint[])
|
||||
UNION
|
||||
SELECT chat_id FROM dialogs
|
||||
WHERE account_id = $1 AND chat_id = ANY($2::bigint[])
|
||||
UNION
|
||||
SELECT scope_id AS chat_id FROM capture_policy
|
||||
WHERE account_id = $1 AND scope_type = 'chat'
|
||||
AND scope_id = ANY($2::bigint[])
|
||||
"""
|
||||
|
||||
_CHAT_ROWS = """
|
||||
WITH ids AS ({ids})
|
||||
SELECT ids.chat_id,
|
||||
@@ -171,6 +183,15 @@ async def get_chat(
|
||||
return _chat_item(row) if row is not None else None
|
||||
|
||||
|
||||
async def get_chats(
|
||||
pool: asyncpg.Pool, account_id: int, ids: list[int]
|
||||
) -> list[ChatListItem]:
|
||||
if not ids:
|
||||
return []
|
||||
rows = await pool.fetch(_CHAT_ROWS.format(ids=_MANY_IDS), account_id, ids)
|
||||
return [_chat_item(row) for row in rows]
|
||||
|
||||
|
||||
async def get_chat_history( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
@@ -308,7 +329,7 @@ async def get_message_versions(
|
||||
rows = await pool.fetch(
|
||||
"SELECT observed_at, edit_date, text FROM message_versions "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"ORDER BY observed_at",
|
||||
"ORDER BY observed_at LIMIT 500",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
|
||||
@@ -54,9 +54,8 @@ WITH chat_meta AS (
|
||||
bool_or(is_broadcast) AS is_broadcast
|
||||
FROM hits GROUP BY chat_id
|
||||
), counts AS (
|
||||
SELECT chat_id, count(*) AS message_count FROM messages
|
||||
SELECT chat_id, message_count FROM chat_stats
|
||||
WHERE account_id = $1 AND chat_id IN (SELECT chat_id FROM merged)
|
||||
GROUP BY chat_id
|
||||
)
|
||||
SELECT m.chat_id, m.title, m.username, m.is_bot, m.is_contact, m.is_broadcast,
|
||||
COALESCE(c.message_count, 0) AS message_count,
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import json
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.files import media_file_name
|
||||
from utils.read.message_view import load_raw
|
||||
from utils.read.models import MediaVersionView, MediaView
|
||||
|
||||
MEDIA_COLS = (
|
||||
@@ -40,15 +41,15 @@ async def _web_page_media_stub(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int, message_id: int
|
||||
) -> MediaView | None:
|
||||
row = await pool.fetchrow(
|
||||
"SELECT date, raw FROM messages "
|
||||
"SELECT date, raw->'web_page' AS web_page FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
)
|
||||
if row is None:
|
||||
if row is None or row["web_page"] is None:
|
||||
return None
|
||||
web_page = load_raw(row["raw"]).get("web_page")
|
||||
web_page = json.loads(row["web_page"])
|
||||
if not isinstance(web_page, dict):
|
||||
return None
|
||||
kind = next(
|
||||
@@ -105,7 +106,7 @@ async def get_media_versions(
|
||||
rows = await pool.fetch(
|
||||
f"SELECT {_VERSION_COLS} FROM media_versions " # noqa: S608
|
||||
"WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 "
|
||||
"ORDER BY observed_at",
|
||||
"ORDER BY observed_at LIMIT 500",
|
||||
account_id,
|
||||
chat_id,
|
||||
message_id,
|
||||
|
||||
@@ -42,7 +42,8 @@ async def get_peer_history(
|
||||
rows = await pool.fetch(
|
||||
"SELECT observed_at, first_name, last_name, username, phone, "
|
||||
"photo_unique_id, is_deleted_account FROM peer_history "
|
||||
"WHERE account_id = $1 AND peer_id = $2 ORDER BY observed_at DESC",
|
||||
"WHERE account_id = $1 AND peer_id = $2 "
|
||||
"ORDER BY observed_at DESC LIMIT 500",
|
||||
account_id,
|
||||
peer_id,
|
||||
)
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
from datetime import datetime
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import asyncpg
|
||||
|
||||
from utils.read.models import Page, PresenceHourly, PresenceSample
|
||||
|
||||
HOURLY_DEFAULT_DAYS = 90
|
||||
|
||||
|
||||
async def presence_history( # noqa: PLR0913
|
||||
pool: asyncpg.Pool,
|
||||
@@ -54,6 +56,8 @@ async def presence_hourly(
|
||||
date_from: datetime | None = None,
|
||||
date_to: datetime | None = None,
|
||||
) -> list[PresenceHourly]:
|
||||
if date_from is None:
|
||||
date_from = datetime.now(UTC) - timedelta(days=HOURLY_DEFAULT_DAYS)
|
||||
params: list[object] = [account_id, peer_id]
|
||||
where = "account_id = $1 AND peer_id = $2"
|
||||
if date_from is not None:
|
||||
|
||||
@@ -42,17 +42,27 @@ async def chat_links(
|
||||
|
||||
|
||||
async def daily_counts(
|
||||
pool: asyncpg.Pool, account_id: int, chat_id: int
|
||||
pool: asyncpg.Pool,
|
||||
account_id: int,
|
||||
chat_id: int,
|
||||
*,
|
||||
date_from: datetime | None = None,
|
||||
date_to: datetime | None = None,
|
||||
) -> list[DayCount]:
|
||||
self_id = await self_user_id(pool, account_id)
|
||||
params: list[object] = [account_id, chat_id, self_id]
|
||||
where = "account_id = $1 AND chat_id = $2"
|
||||
if date_from is not None:
|
||||
params.append(date_from)
|
||||
where += f" AND date >= ${len(params)}"
|
||||
if date_to is not None:
|
||||
params.append(date_to)
|
||||
where += f" AND date < ${len(params)}"
|
||||
rows = await pool.fetch(
|
||||
"SELECT date_trunc('day', date) AS day, count(*) AS count, "
|
||||
"SELECT date_trunc('day', date) AS day, count(*) AS count, " # noqa: S608
|
||||
"count(*) FILTER (WHERE sender_id = $3) AS outgoing FROM messages "
|
||||
"WHERE account_id = $1 AND chat_id = $2 "
|
||||
"GROUP BY day ORDER BY day",
|
||||
account_id,
|
||||
chat_id,
|
||||
self_id,
|
||||
f"WHERE {where} GROUP BY day ORDER BY day",
|
||||
*params,
|
||||
)
|
||||
return [DayCount(**dict(row)) for row in rows]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user