feat(backend): storage lists the heaviest chats, summary and file list take a chat filter

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(

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

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:
+def _groups(rows: list[asyncpg.Record]) -> list[StorageGroupView]:
-    rows = await pool.fetch(_SUMMARY, account_id, HEAVY_BYTES)
-    trash = await pool.fetchrow(_TRASH_TOTAL, account_id)
     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"],