From f646f29ff967dc355f793300c5eb99c760e71663 Mon Sep 17 00:00:00 2001 From: h Date: Wed, 14 Oct 2026 04:13:13 +0200 Subject: [PATCH] feat(backend): policy rules show what they inherit, files of a chat or folder go to trash in one batch, trash list is paged --- backend/src/api/routers/policy.py | 30 +++++++- backend/src/api/routers/storage.py | 49 ++++++++++++- backend/src/utils/policy/models.py | 18 +++++ backend/src/utils/policy/overview.py | 103 +++++++++++++++++++++++++++ backend/src/utils/policy/resolver.py | 28 ++++++-- backend/src/utils/read/chats.py | 12 ++++ backend/src/utils/read/models.py | 8 +++ backend/src/utils/read/storage.py | 93 ++++++++++++++++++++++-- backend/src/utils/trash.py | 63 +++++++++++++++- 9 files changed, 389 insertions(+), 15 deletions(-) create mode 100644 backend/src/utils/policy/overview.py diff --git a/backend/src/api/routers/policy.py b/backend/src/api/routers/policy.py index 4241194..84abec1 100644 --- a/backend/src/api/routers/policy.py +++ b/backend/src/api/routers/policy.py @@ -5,13 +5,15 @@ from dishka.integrations.fastapi import DishkaRoute, FromDishka from fastapi import APIRouter, HTTPException, Query from pydantic import BaseModel -from utils.policy import repository +from utils.policy import overview, repository from utils.policy.models import ( CaptureToggles, ChatKind, ChatMeta, PolicyCreate, PolicyRecord, + PolicyRuleView, + PolicyScopeView, ) from utils.policy.resolver import resolve @@ -50,6 +52,32 @@ async def effective_policy( return resolve(chat, folders, policies) +@router.get("/overview") +async def policy_overview( + pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()] +) -> list[PolicyRuleView]: + return await overview.list_rules(pool, account_id) + + +@router.get("/scope") +async def policy_scope( + pool: FromDishka[asyncpg.Pool], + account_id: Annotated[int, Query()], + chat_id: Annotated[int | None, Query()] = None, + folder_id: Annotated[int | None, Query()] = None, +) -> PolicyScopeView: + if chat_id is not None: + return await overview.chat_scope(pool, account_id, chat_id) + scope = ( + await overview.folder_scope(pool, account_id, folder_id) + if folder_id is not None + else None + ) + if scope is None: + raise HTTPException(status_code=404, detail="scope not found") + return scope + + @router.post("", status_code=201) async def create_policy( pool: FromDishka[asyncpg.Pool], body: PolicyCreate diff --git a/backend/src/api/routers/storage.py b/backend/src/api/routers/storage.py index 45911b2..e487c44 100644 --- a/backend/src/api/routers/storage.py +++ b/backend/src/api/routers/storage.py @@ -14,6 +14,7 @@ from utils.read.models import ( DEFAULT_LIMIT, Page, StorageFileView, + StorageScopeView, StorageSummaryView, TrashItemView, TrashOutcome, @@ -30,7 +31,13 @@ class TrashRequest(BaseModel): class RestoreRequest(BaseModel): account_id: int - keys: list[str] + keys: list[str] | None = None + + +class ScopeRequest(BaseModel): + account_id: int + chat_id: int | None = None + folder_id: int | None = None class PurgeRequest(BaseModel): @@ -57,11 +64,47 @@ async def files( return await storage_read.list_files(pool, account_id, group, page) +async def _chat_ids( + pool: asyncpg.Pool, account_id: int, chat_id: int | None, folder_id: int | None +) -> list[int]: + chat_ids = await storage_read.scope_chat_ids(pool, account_id, chat_id, folder_id) + if chat_ids is None: + raise HTTPException(status_code=404, detail="scope not found") + return chat_ids + + +@router.get("/scope") +async def scope( + pool: FromDishka[asyncpg.Pool], + account_id: Annotated[int, Query()], + chat_id: Annotated[int | None, Query()] = None, + folder_id: Annotated[int | None, Query()] = None, +) -> StorageScopeView: + chat_ids = await _chat_ids(pool, account_id, chat_id, folder_id) + return await storage_read.get_scope(pool, account_id, chat_ids) + + @router.get("/trash") async def trash_items( - pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()] + pool: FromDishka[asyncpg.Pool], + account_id: Annotated[int, Query()], + limit: Annotated[int, Query()] = DEFAULT_LIMIT, + offset: Annotated[int, Query()] = 0, ) -> list[TrashItemView]: - return await storage_read.list_trash(pool, account_id) + page = Page(limit=limit, offset=offset) + return await storage_read.list_trash(pool, account_id, page) + + +@router.post("/trash/scope") +async def trash_scope( + pool: FromDishka[asyncpg.Pool], + storage: FromDishka[ContentAddressedStorage], + body: ScopeRequest, +) -> TrashOutcome: + chat_ids = await _chat_ids(pool, body.account_id, body.chat_id, body.folder_id) + blobs = await pool.fetch(storage_read.SCOPE_BLOBS, body.account_id, chat_ids) + free = [blob for blob in blobs if blob["free"]] + return await trash.trash_blobs(pool, storage, body.account_id, free) @router.post("/trash") diff --git a/backend/src/utils/policy/models.py b/backend/src/utils/policy/models.py index a2dce6e..3e5a3d1 100644 --- a/backend/src/utils/policy/models.py +++ b/backend/src/utils/policy/models.py @@ -79,6 +79,24 @@ class FolderSpec: is_chatlist: bool = False +class PolicyParent(BaseModel): + scope_type: ScopeType + scope_id: int | None + title: str | None = None + toggles: CaptureToggles + + +class PolicyRuleView(PolicyRecord): + title: str | None = None + parent: PolicyParent | None = None + + +class PolicyScopeView(BaseModel): + title: str | None + rule: PolicyRecord | None + parent: PolicyParent | None + + @dataclass(slots=True) class PolicySet: chat: dict[int, CaptureToggles] = field(default_factory=dict) diff --git a/backend/src/utils/policy/overview.py b/backend/src/utils/policy/overview.py new file mode 100644 index 0000000..a5a05b7 --- /dev/null +++ b/backend/src/utils/policy/overview.py @@ -0,0 +1,103 @@ +import asyncpg + +from utils.policy import repository +from utils.policy.models import ( + ChatKind, + ChatMeta, + FolderSpec, + PolicyParent, + PolicyRecord, + PolicyRuleView, + PolicyScopeView, + PolicySet, + ScopeType, +) +from utils.policy.resolver import resolve_parent +from utils.read import chats +from utils.read.models import ChatListItem + + +def _meta(chat_id: int, chat: ChatListItem | None) -> ChatMeta: + if chat_id > 0: + kind = ChatKind.DM + elif chat is not None and chat.is_broadcast: + kind = ChatKind.CHANNEL + else: + kind = ChatKind.GROUP + return ChatMeta( + chat_id=chat_id, + kind=kind, + is_bot=bool(chat and chat.is_bot), + is_contact=chat.is_contact if chat else None, + ) + + +def _without_chats(policies: PolicySet) -> PolicySet: + return PolicySet(folder=policies.folder, defaults=policies.defaults) + + +def _chat_parent( + chat_id: int, + chat: ChatListItem | None, + folders: list[FolderSpec], + policies: PolicySet, +) -> PolicyParent: + return resolve_parent(_meta(chat_id, chat), folders, _without_chats(policies)) + + +async def list_rules(pool: asyncpg.Pool, account_id: int) -> list[PolicyRuleView]: + records = await repository.list_policies(pool, account_id) + folders = await repository.list_folders(pool, account_id) + policies = await repository.load_policy_set(pool, account_id) + chat_ids = [ + record.scope_id + for record in records + if record.scope_type is ScopeType.CHAT and record.scope_id is not None + ] + known = { + chat.chat_id: chat for chat in await chats.get_chats(pool, account_id, chat_ids) + } + titles = {folder.folder_id: folder.title for folder in folders} + rules: list[PolicyRuleView] = [] + for record in records: + rule = PolicyRuleView(**record.model_dump()) + if record.scope_type is ScopeType.FOLDER and record.scope_id is not None: + rule.title = titles.get(record.scope_id) + if record.scope_type is ScopeType.CHAT and record.scope_id is not None: + chat = known.get(record.scope_id) + rule.title = chat.title if chat else None + rule.parent = _chat_parent(record.scope_id, chat, folders, policies) + rules.append(rule) + return rules + + +async def _rule( + pool: asyncpg.Pool, account_id: int, scope_type: ScopeType, scope_id: int +) -> PolicyRecord | None: + return await repository.find_policy(pool, account_id, scope_type, scope_id) + + +async def chat_scope( + pool: asyncpg.Pool, account_id: int, chat_id: int +) -> PolicyScopeView: + folders = await repository.list_folders(pool, account_id) + policies = await repository.load_policy_set(pool, account_id) + chat = await chats.get_chat(pool, account_id, chat_id) + return PolicyScopeView( + title=chat.title if chat else None, + rule=await _rule(pool, account_id, ScopeType.CHAT, chat_id), + parent=_chat_parent(chat_id, chat, folders, policies), + ) + + +async def folder_scope( + pool: asyncpg.Pool, account_id: int, folder_id: int +) -> PolicyScopeView | None: + folder = await repository.get_folder(pool, account_id, folder_id) + if folder is None: + return None + return PolicyScopeView( + title=folder.title, + rule=await _rule(pool, account_id, ScopeType.FOLDER, folder_id), + parent=None, + ) diff --git a/backend/src/utils/policy/resolver.py b/backend/src/utils/policy/resolver.py index 7b597f1..fdffaef 100644 --- a/backend/src/utils/policy/resolver.py +++ b/backend/src/utils/policy/resolver.py @@ -1,11 +1,14 @@ from collections.abc import Iterable from utils.policy.models import ( + DEFAULT_SCOPE_BY_KIND, CaptureToggles, ChatKind, ChatMeta, FolderSpec, + PolicyParent, PolicySet, + ScopeType, ) @@ -33,14 +36,29 @@ def folder_contains(folder: FolderSpec, chat: ChatMeta) -> bool: return _category_match(folder, chat) +def resolve_parent( + chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet +) -> PolicyParent: + for folder in sorted(folders, key=lambda f: f.order_index): + toggles = policies.folder.get(folder.folder_id) + if toggles is not None and folder_contains(folder, chat): + return PolicyParent( + scope_type=ScopeType.FOLDER, + scope_id=folder.folder_id, + title=folder.title, + toggles=toggles, + ) + return PolicyParent( + scope_type=DEFAULT_SCOPE_BY_KIND[chat.kind], + scope_id=None, + toggles=policies.defaults.get(chat.kind, CaptureToggles()), + ) + + def resolve( chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet ) -> CaptureToggles: chat_override = policies.chat.get(chat.chat_id) if chat_override is not None: return chat_override - for folder in sorted(folders, key=lambda f: f.order_index): - toggles = policies.folder.get(folder.folder_id) - if toggles is not None and folder_contains(folder, chat): - return toggles - return policies.defaults.get(chat.kind, CaptureToggles()) + return resolve_parent(chat, folders, policies).toggles diff --git a/backend/src/utils/read/chats.py b/backend/src/utils/read/chats.py index d54ecf4..dca0d53 100644 --- a/backend/src/utils/read/chats.py +++ b/backend/src/utils/read/chats.py @@ -176,6 +176,18 @@ async def list_chats( return [_chat_item(row) for row in rows] +async def folder_chat_ids( + pool: asyncpg.Pool, account_id: int, folder: FolderSpec +) -> list[int]: + rows_sql = _CHAT_ROWS.format(ids=_ALL_IDS) + rows = await pool.fetch( + f"SELECT chat.chat_id FROM ({rows_sql}) chat WHERE {_folder_filter(1)}", # noqa: S608 + account_id, + *_folder_params(folder), + ) + return [row["chat_id"] for row in rows] + + async def get_chat( pool: asyncpg.Pool, account_id: int, chat_id: int ) -> ChatListItem | None: diff --git a/backend/src/utils/read/models.py b/backend/src/utils/read/models.py index 2ec8ccf..d323e8b 100644 --- a/backend/src/utils/read/models.py +++ b/backend/src/utils/read/models.py @@ -287,6 +287,14 @@ class TrashOutcome(BaseModel): blocked: list[int] = [] +class StorageScopeView(BaseModel): + chats: int + files: int + bytes: int + kept: int + gone: int + + class AvatarRef(BaseModel): unique_id: str storage_key: str | None diff --git a/backend/src/utils/read/storage.py b/backend/src/utils/read/storage.py index 736d807..625bc8a 100644 --- a/backend/src/utils/read/storage.py +++ b/backend/src/utils/read/storage.py @@ -4,10 +4,13 @@ from collections import defaultdict import asyncpg from utils.files import media_file_name +from utils.policy import repository as policy_repository +from utils.read import chats from utils.read.models import ( Page, StorageFileView, StorageGroupView, + StorageScopeView, StorageSummaryView, TrashItemView, ) @@ -71,6 +74,46 @@ SELECT storage_key, chat_id, message_id, kind, mime, file_name, file_size, trashed_at FROM media_trash WHERE account_id = $1 ORDER BY file_size DESC NULLS LAST, storage_key +LIMIT $2 OFFSET $3 +""" + +SCOPE_BLOBS = """ +WITH inside AS ( + SELECT DISTINCT storage_key FROM media + WHERE account_id = $1 AND chat_id = ANY($2::bigint[]) + AND downloaded AND storage_key IS NOT NULL +), blobs AS ( + SELECT m.storage_key, max(m.file_size) AS file_size, + min(m.id) FILTER ( + WHERE m.account_id = $1 AND m.chat_id = ANY($2::bigint[]) + ) AS media_id, + bool_and(m.account_id = $1 AND m.chat_id = ANY($2::bigint[])) AS own + FROM media m JOIN inside i ON i.storage_key = m.storage_key + GROUP BY m.storage_key +) +SELECT b.storage_key, b.file_size, r.chat_id, r.message_id, r.kind, r.mime, + b.own + AND NOT EXISTS (SELECT 1 FROM stories s WHERE s.storage_key = b.storage_key) + AND NOT EXISTS (SELECT 1 FROM avatars a WHERE a.storage_key = b.storage_key) + AND NOT EXISTS ( + SELECT 1 FROM custom_emoji e WHERE e.storage_key = b.storage_key + ) AS free +FROM blobs b JOIN media r ON r.id = b.media_id +""" + +_SCOPE_GONE = """ +SELECT ( + SELECT count(DISTINCT m.storage_key) FROM media m + JOIN messages g ON g.account_id = m.account_id AND g.chat_id = m.chat_id + AND g.message_id = m.message_id + WHERE m.account_id = $1 AND m.chat_id = ANY($2::bigint[]) + AND m.downloaded AND m.storage_key IS NOT NULL + AND g.deleted_at IS NOT NULL +) + ( + SELECT count(DISTINCT storage_key) FROM media + WHERE account_id = $1 AND chat_id = ANY($2::bigint[]) + AND downloaded AND storage_key IS NOT NULL AND ttl_seconds IS NOT NULL +) """ _TRASH_REFS = """ @@ -179,20 +222,26 @@ async def list_files( return list(files.values()) -async def list_trash(pool: asyncpg.Pool, account_id: int) -> list[TrashItemView]: - rows = await pool.fetch(_TRASH, account_id) +async def list_trash( + pool: asyncpg.Pool, account_id: int, page: Page +) -> list[TrashItemView]: + rows = await pool.fetch(_TRASH, account_id, page.capped_limit, page.offset) refs = await pool.fetch( _TRASH_REFS, account_id, [row["storage_key"] for row in rows], REFS_CHECKED ) messages = await _messages(pool, account_id, refs) alive: dict[str, bool] = defaultdict(bool) counts: dict[str, int] = {} + names: dict[str, str] = {} for ref, message in zip(refs, messages, strict=True): - alive[ref["storage_key"]] |= _alive(ref, message) - counts[ref["storage_key"]] = ref["refs"] + key = ref["storage_key"] + alive[key] |= _alive(ref, message) + counts[key] = ref["refs"] + if message is not None and message["original_name"]: + names.setdefault(key, message["original_name"]) return [ TrashItemView( - **dict(row), + **{**dict(row), "file_name": _trash_name(row, names)}, purge_at=row["trashed_at"] + RETENTION, refs=counts.get(row["storage_key"], 0), refetchable=alive.get(row["storage_key"], True), @@ -201,6 +250,40 @@ async def list_trash(pool: asyncpg.Pool, account_id: int) -> list[TrashItemView] ] +def _trash_name(row: asyncpg.Record, names: dict[str, str]) -> str: + original = names.get(row["storage_key"]) + if original is None: + return row["file_name"] + return media_file_name(row["kind"], row["mime"], row["message_id"], original) + + +async def scope_chat_ids( + pool: asyncpg.Pool, account_id: int, chat_id: int | None, folder_id: int | None +) -> list[int] | None: + if chat_id is not None: + return [chat_id] + if folder_id is None: + return None + folder = await policy_repository.get_folder(pool, account_id, folder_id) + if folder is None: + return None + return await chats.folder_chat_ids(pool, account_id, folder) + + +async def get_scope( + pool: asyncpg.Pool, account_id: int, chat_ids: list[int] +) -> StorageScopeView: + blobs = await pool.fetch(SCOPE_BLOBS, account_id, chat_ids) + free = [blob for blob in blobs if blob["free"]] + return StorageScopeView( + chats=len(chat_ids), + files=len(free), + bytes=sum(blob["file_size"] or 0 for blob in free), + kept=len(blobs) - len(free), + gone=await pool.fetchval(_SCOPE_GONE, account_id, chat_ids) if free else 0, + ) + + async def get_trash_item(pool: asyncpg.Pool, storage_key: str) -> asyncpg.Record | None: return await pool.fetchrow( "SELECT kind, mime, file_name FROM media_trash WHERE storage_key = $1", diff --git a/backend/src/utils/trash.py b/backend/src/utils/trash.py index 30d30f5..3a973d0 100644 --- a/backend/src/utils/trash.py +++ b/backend/src/utils/trash.py @@ -4,6 +4,7 @@ from datetime import timedelta import asyncpg +from utils.files import media_file_name from utils.read.media import get_media from utils.read.models import TrashOutcome from utils.storage import ContentAddressedStorage @@ -12,6 +13,7 @@ logger = logging.getLogger(__name__) RETENTION = timedelta(days=7) PURGE_INTERVAL_SECONDS = 3600 +BATCH = 200 _SHARED_ELSEWHERE = """ SELECT EXISTS (SELECT 1 FROM stories WHERE storage_key = $1) @@ -47,6 +49,12 @@ _FORGET_KEYS = """ UPDATE media SET evicted_key = NULL WHERE evicted_key = ANY($1::text[]) """ +_EVICT_MANY = """ +UPDATE media SET evicted_key = storage_key, storage_key = NULL, + downloaded = false, evicted_at = now() +WHERE storage_key = ANY($1::text[]) +""" + async def move_to_trash( pool: asyncpg.Pool, @@ -86,12 +94,65 @@ async def move_to_trash( return outcome +def _move_batch(storage: ContentAddressedStorage, keys: list[str]) -> None: + moved: list[str] = [] + try: + for key in keys: + if storage.exists(key): + storage.trash(key) + moved.append(key) + except OSError: + for key in moved: + storage.restore(key) + raise + + +async def trash_blobs( + pool: asyncpg.Pool, + storage: ContentAddressedStorage, + account_id: int, + blobs: list[asyncpg.Record], +) -> TrashOutcome: + outcome = TrashOutcome() + for start in range(0, len(blobs), BATCH): + batch = blobs[start : start + BATCH] + keys = [blob["storage_key"] for blob in batch] + async with pool.acquire() as conn, conn.transaction(): + await conn.execute(_EVICT_MANY, keys) + await conn.executemany( + _INSERT, + [ + ( + blob["storage_key"], + account_id, + blob["chat_id"], + blob["message_id"], + blob["kind"], + blob["mime"], + media_file_name(blob["kind"], blob["mime"], blob["message_id"]), + blob["file_size"], + ) + for blob in batch + ], + ) + await asyncio.to_thread(_move_batch, storage, keys) + outcome.keys.extend(keys) + outcome.files += len(batch) + outcome.bytes += sum(blob["file_size"] or 0 for blob in batch) + return outcome + + async def restore( pool: asyncpg.Pool, storage: ContentAddressedStorage, account_id: int, - keys: list[str], + keys: list[str] | None, ) -> TrashOutcome: + if keys is None: + rows = await pool.fetch( + "SELECT storage_key FROM media_trash WHERE account_id = $1", account_id + ) + keys = [row["storage_key"] for row in rows] outcome = TrashOutcome() for key in keys: async with pool.acquire() as conn, conn.transaction():