From 5eacabd10196abf91454cd6afaa117cb83697ba1 Mon Sep 17 00:00:00 2001 From: h Date: Wed, 14 Oct 2026 05:02:03 +0200 Subject: [PATCH] feat(backend): storage lists the heaviest chats, summary and file list take a chat filter --- backend/src/api/routers/storage.py | 22 +++++++- backend/src/utils/read/models.py | 7 +++ backend/src/utils/read/storage.py | 91 ++++++++++++++++++++++++++---- 3 files changed, 106 insertions(+), 14 deletions(-) diff --git a/backend/src/api/routers/storage.py b/backend/src/api/routers/storage.py index e487c44..a7f1513 100644 --- a/backend/src/api/routers/storage.py +++ b/backend/src/api/routers/storage.py @@ -13,6 +13,7 @@ from utils.read import storage as storage_read from utils.read.models import ( DEFAULT_LIMIT, Page, + StorageChatView, StorageFileView, StorageScopeView, StorageSummaryView, @@ -47,9 +48,11 @@ class PurgeRequest(BaseModel): @router.get("/summary") async def summary( - pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()] + pool: FromDishka[asyncpg.Pool], + account_id: Annotated[int, Query()], + chat_id: Annotated[int | None, Query()] = None, ) -> StorageSummaryView: - return await storage_read.get_summary(pool, account_id) + return await storage_read.get_summary(pool, account_id, chat_id) @router.get("/files") @@ -57,11 +60,24 @@ async def files( pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()], group: Annotated[str | None, Query()] = None, + chat_id: Annotated[int | None, Query()] = None, limit: Annotated[int, Query()] = DEFAULT_LIMIT, offset: Annotated[int, Query()] = 0, ) -> list[StorageFileView]: page = Page(limit=limit, offset=offset) - return await storage_read.list_files(pool, account_id, group, page) + return await storage_read.list_files(pool, account_id, group, page, chat_id) + + +@router.get("/chats") +async def chats( + pool: FromDishka[asyncpg.Pool], + account_id: Annotated[int, Query()], + group: Annotated[str | None, Query()] = None, + limit: Annotated[int, Query()] = DEFAULT_LIMIT, + offset: Annotated[int, Query()] = 0, +) -> list[StorageChatView]: + page = Page(limit=limit, offset=offset) + return await storage_read.list_chats(pool, account_id, group, page) async def _chat_ids( diff --git a/backend/src/utils/read/models.py b/backend/src/utils/read/models.py index d323e8b..8cc6f50 100644 --- a/backend/src/utils/read/models.py +++ b/backend/src/utils/read/models.py @@ -253,6 +253,13 @@ class StorageSummaryView(BaseModel): groups: list[StorageGroupView] +class StorageChatView(BaseModel): + chat_id: int + files: int + bytes: int + groups: list[StorageGroupView] + + class StorageFileView(BaseModel): media_id: int chat_id: int diff --git a/backend/src/utils/read/storage.py b/backend/src/utils/read/storage.py index 625bc8a..3406e3d 100644 --- a/backend/src/utils/read/storage.py +++ b/backend/src/utils/read/storage.py @@ -1,5 +1,7 @@ import asyncio from collections import defaultdict +from itertools import groupby +from operator import itemgetter import asyncpg @@ -8,6 +10,7 @@ from utils.policy import repository as policy_repository from utils.read import chats from utils.read.models import ( Page, + StorageChatView, StorageFileView, StorageGroupView, StorageScopeView, @@ -34,6 +37,7 @@ WITH blobs AS ( SELECT storage_key, max(file_size) AS size, min(kind) AS kind FROM media WHERE account_id = $1 AND downloaded AND storage_key IS NOT NULL + AND ($3::bigint IS NULL OR chat_id = $3) GROUP BY storage_key ) SELECT kind, count(*) AS files, coalesce(sum(size), 0)::bigint AS bytes, @@ -49,24 +53,49 @@ FROM media_trash WHERE account_id = $1 _FILES = """ WITH top AS MATERIALIZED ( - SELECT storage_key, max(file_size) AS file_size, count(*) AS refs + SELECT storage_key, max(file_size) AS file_size FROM media WHERE account_id = $1 AND downloaded AND storage_key IS NOT NULL AND ($2::text[] IS NULL OR (kind = ANY($2)) = $3) + AND ($7::bigint IS NULL OR chat_id = $7) GROUP BY storage_key ORDER BY 2 DESC NULLS LAST, 1 LIMIT $4 OFFSET $5 ) SELECT * FROM ( - SELECT t.storage_key, t.file_size, t.refs, m.id, m.chat_id, m.message_id, + SELECT t.storage_key, t.file_size, m.id, m.chat_id, m.message_id, m.kind, m.mime, m.ttl_seconds, - row_number() OVER (PARTITION BY t.storage_key ORDER BY m.id) AS n + count(*) OVER (PARTITION BY t.storage_key) AS refs, + row_number() OVER ( + PARTITION BY t.storage_key + ORDER BY (m.chat_id = $7) DESC NULLS LAST, m.id + ) AS n FROM top t JOIN media m ON m.storage_key = t.storage_key AND m.account_id = $1 AND m.downloaded ) ranked WHERE n <= $6 -ORDER BY file_size DESC NULLS LAST, storage_key, id +ORDER BY file_size DESC NULLS LAST, storage_key, n +""" + +_CHATS = """ +WITH blobs AS ( + SELECT chat_id, storage_key, max(file_size) AS size, min(kind) AS kind + FROM media + WHERE account_id = $1 AND downloaded AND storage_key IS NOT NULL + AND ($2::text[] IS NULL OR (kind = ANY($2)) = $3) + GROUP BY chat_id, storage_key +), top AS ( + SELECT chat_id, coalesce(sum(size), 0)::bigint AS total + FROM blobs GROUP BY chat_id + ORDER BY 2 DESC, 1 + LIMIT $4 OFFSET $5 +) +SELECT t.chat_id, t.total, b.kind, count(*) AS files, + coalesce(sum(b.size), 0)::bigint AS bytes +FROM top t JOIN blobs b USING (chat_id) +GROUP BY t.chat_id, t.total, b.kind +ORDER BY t.total DESC, t.chat_id """ _TRASH = """ @@ -143,9 +172,7 @@ def _group_filter(group: str | None) -> tuple[list[str] | None, bool]: return list(GROUP_KINDS.get(group, ())), True -async def get_summary(pool: asyncpg.Pool, account_id: int) -> StorageSummaryView: - rows = await pool.fetch(_SUMMARY, account_id, HEAVY_BYTES) - trash = await pool.fetchrow(_TRASH_TOTAL, account_id) +def _groups(rows: list[asyncpg.Record]) -> list[StorageGroupView]: groups = { name: StorageGroupView(group=name, files=0, bytes=0) for name in (*GROUP_KINDS, OTHER_GROUP) @@ -154,6 +181,14 @@ async def get_summary(pool: asyncpg.Pool, account_id: int) -> StorageSummaryView group = groups[_KIND_GROUP.get(row["kind"], OTHER_GROUP)] group.files += row["files"] group.bytes += row["bytes"] + return list(groups.values()) + + +async def get_summary( + pool: asyncpg.Pool, account_id: int, chat_id: int | None = None +) -> StorageSummaryView: + rows = await pool.fetch(_SUMMARY, account_id, HEAVY_BYTES, chat_id) + trash = await pool.fetchrow(_TRASH_TOTAL, account_id) return StorageSummaryView( files=sum(row["files"] for row in rows), bytes=sum(row["bytes"] for row in rows), @@ -162,10 +197,31 @@ async def get_summary(pool: asyncpg.Pool, account_id: int) -> StorageSummaryView heavy_threshold=HEAVY_BYTES, trash_files=trash["files"], trash_bytes=trash["bytes"], - groups=list(groups.values()), + groups=_groups(rows), ) +async def list_chats( + pool: asyncpg.Pool, account_id: int, group: str | None, page: Page +) -> list[StorageChatView]: + kinds, include = _group_filter(group) + rows = await pool.fetch( + _CHATS, account_id, kinds, include, page.capped_limit, page.offset + ) + return [ + StorageChatView( + chat_id=chat_id, + files=sum(row["files"] for row in items), + bytes=items[0]["total"], + groups=_groups(items), + ) + for chat_id, items in ( + (chat_id, list(chunk)) + for chat_id, chunk in groupby(rows, key=itemgetter("chat_id")) + ) + ] + + async def _messages( pool: asyncpg.Pool, account_id: int, refs: list[asyncpg.Record] ) -> list[asyncpg.Record | None]: @@ -189,18 +245,31 @@ def _alive(ref: asyncpg.Record, message: asyncpg.Record | None) -> bool: async def list_files( - pool: asyncpg.Pool, account_id: int, group: str | None, page: Page + pool: asyncpg.Pool, + account_id: int, + group: str | None, + page: Page, + chat_id: int | None = None, ) -> list[StorageFileView]: kinds, include = _group_filter(group) refs = await pool.fetch( - _FILES, account_id, kinds, include, page.capped_limit, page.offset, REFS_CHECKED + _FILES, + account_id, + kinds, + include, + page.capped_limit, + page.offset, + REFS_CHECKED, + chat_id, ) messages = await _messages(pool, account_id, refs) files: dict[str, StorageFileView] = {} for ref, message in zip(refs, messages, strict=True): alive = _alive(ref, message) known = files.get(ref["storage_key"]) - if known is not None and (known.refetchable or not alive): + inside = chat_id is None or ref["chat_id"] == chat_id + if known is not None and not (alive and inside and not known.refetchable): + known.refetchable = known.refetchable or alive continue files[ref["storage_key"]] = StorageFileView( media_id=ref["id"],