feat(backend): storage lists the heaviest chats, summary and file list take a chat filter
This commit is contained in:
1 parent
0362777a51
commit
5eacabd101
3 files changed
+106
-14
No files matched your search
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"],
|
||||
|
||||
Reference in new issue
Block a user