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

This commit is contained in:
hh committed 2026-10-14 05:02:03 +02:00
1 parent 0362777a51
commit 5eacabd101
3 files changed
+106 -14

No files matched your search

+19 -3
View File
@@ -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(
+7
View File
@@ -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
+80 -11
View File
@@ -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"],