feat(backend): policy rules show what they inherit, files of a chat or folder go to trash in one batch, trash list is paged

This commit is contained in:
hh committed 2026-10-14 04:13:13 +02:00
1 parent 945adcf879
commit f646f29ff9
9 files changed
+389 -15

No files matched your search

+29 -1
View File
@@ -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
+46 -3
View File
@@ -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")
+18
View File
@@ -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)
+103
View File
@@ -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,
)
+23 -5
View File
@@ -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
+12
View File
@@ -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:
+8
View File
@@ -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
+88 -5
View File
@@ -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",
+62 -1
View File
@@ -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():