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 fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel from pydantic import BaseModel
from utils.policy import repository from utils.policy import overview, repository
from utils.policy.models import ( from utils.policy.models import (
CaptureToggles, CaptureToggles,
ChatKind, ChatKind,
ChatMeta, ChatMeta,
PolicyCreate, PolicyCreate,
PolicyRecord, PolicyRecord,
PolicyRuleView,
PolicyScopeView,
) )
from utils.policy.resolver import resolve from utils.policy.resolver import resolve
@@ -50,6 +52,32 @@ async def effective_policy(
return resolve(chat, folders, policies) 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) @router.post("", status_code=201)
async def create_policy( async def create_policy(
pool: FromDishka[asyncpg.Pool], body: PolicyCreate pool: FromDishka[asyncpg.Pool], body: PolicyCreate
+46 -3
View File
@@ -14,6 +14,7 @@ from utils.read.models import (
DEFAULT_LIMIT, DEFAULT_LIMIT,
Page, Page,
StorageFileView, StorageFileView,
StorageScopeView,
StorageSummaryView, StorageSummaryView,
TrashItemView, TrashItemView,
TrashOutcome, TrashOutcome,
@@ -30,7 +31,13 @@ class TrashRequest(BaseModel):
class RestoreRequest(BaseModel): class RestoreRequest(BaseModel):
account_id: int 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): class PurgeRequest(BaseModel):
@@ -57,11 +64,47 @@ async def files(
return await storage_read.list_files(pool, account_id, group, page) 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") @router.get("/trash")
async def trash_items( 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]: ) -> 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") @router.post("/trash")
+18
View File
@@ -79,6 +79,24 @@ class FolderSpec:
is_chatlist: bool = False 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) @dataclass(slots=True)
class PolicySet: class PolicySet:
chat: dict[int, CaptureToggles] = field(default_factory=dict) 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 collections.abc import Iterable
from utils.policy.models import ( from utils.policy.models import (
DEFAULT_SCOPE_BY_KIND,
CaptureToggles, CaptureToggles,
ChatKind, ChatKind,
ChatMeta, ChatMeta,
FolderSpec, FolderSpec,
PolicyParent,
PolicySet, PolicySet,
ScopeType,
) )
@@ -33,14 +36,29 @@ def folder_contains(folder: FolderSpec, chat: ChatMeta) -> bool:
return _category_match(folder, chat) 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( def resolve(
chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet chat: ChatMeta, folders: Iterable[FolderSpec], policies: PolicySet
) -> CaptureToggles: ) -> CaptureToggles:
chat_override = policies.chat.get(chat.chat_id) chat_override = policies.chat.get(chat.chat_id)
if chat_override is not None: if chat_override is not None:
return chat_override return chat_override
for folder in sorted(folders, key=lambda f: f.order_index): return resolve_parent(chat, folders, policies).toggles
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())
+12
View File
@@ -176,6 +176,18 @@ async def list_chats(
return [_chat_item(row) for row in rows] 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( async def get_chat(
pool: asyncpg.Pool, account_id: int, chat_id: int pool: asyncpg.Pool, account_id: int, chat_id: int
) -> ChatListItem | None: ) -> ChatListItem | None:
+8
View File
@@ -287,6 +287,14 @@ class TrashOutcome(BaseModel):
blocked: list[int] = [] blocked: list[int] = []
class StorageScopeView(BaseModel):
chats: int
files: int
bytes: int
kept: int
gone: int
class AvatarRef(BaseModel): class AvatarRef(BaseModel):
unique_id: str unique_id: str
storage_key: str | None storage_key: str | None
+88 -5
View File
@@ -4,10 +4,13 @@ from collections import defaultdict
import asyncpg import asyncpg
from utils.files import media_file_name 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 ( from utils.read.models import (
Page, Page,
StorageFileView, StorageFileView,
StorageGroupView, StorageGroupView,
StorageScopeView,
StorageSummaryView, StorageSummaryView,
TrashItemView, TrashItemView,
) )
@@ -71,6 +74,46 @@ SELECT storage_key, chat_id, message_id, kind, mime, file_name, file_size,
trashed_at trashed_at
FROM media_trash WHERE account_id = $1 FROM media_trash WHERE account_id = $1
ORDER BY file_size DESC NULLS LAST, storage_key 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 = """ _TRASH_REFS = """
@@ -179,20 +222,26 @@ async def list_files(
return list(files.values()) return list(files.values())
async def list_trash(pool: asyncpg.Pool, account_id: int) -> list[TrashItemView]: async def list_trash(
rows = await pool.fetch(_TRASH, account_id) 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( refs = await pool.fetch(
_TRASH_REFS, account_id, [row["storage_key"] for row in rows], REFS_CHECKED _TRASH_REFS, account_id, [row["storage_key"] for row in rows], REFS_CHECKED
) )
messages = await _messages(pool, account_id, refs) messages = await _messages(pool, account_id, refs)
alive: dict[str, bool] = defaultdict(bool) alive: dict[str, bool] = defaultdict(bool)
counts: dict[str, int] = {} counts: dict[str, int] = {}
names: dict[str, str] = {}
for ref, message in zip(refs, messages, strict=True): for ref, message in zip(refs, messages, strict=True):
alive[ref["storage_key"]] |= _alive(ref, message) key = ref["storage_key"]
counts[ref["storage_key"]] = ref["refs"] 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 [ return [
TrashItemView( TrashItemView(
**dict(row), **{**dict(row), "file_name": _trash_name(row, names)},
purge_at=row["trashed_at"] + RETENTION, purge_at=row["trashed_at"] + RETENTION,
refs=counts.get(row["storage_key"], 0), refs=counts.get(row["storage_key"], 0),
refetchable=alive.get(row["storage_key"], True), 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: async def get_trash_item(pool: asyncpg.Pool, storage_key: str) -> asyncpg.Record | None:
return await pool.fetchrow( return await pool.fetchrow(
"SELECT kind, mime, file_name FROM media_trash WHERE storage_key = $1", "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 import asyncpg
from utils.files import media_file_name
from utils.read.media import get_media from utils.read.media import get_media
from utils.read.models import TrashOutcome from utils.read.models import TrashOutcome
from utils.storage import ContentAddressedStorage from utils.storage import ContentAddressedStorage
@@ -12,6 +13,7 @@ logger = logging.getLogger(__name__)
RETENTION = timedelta(days=7) RETENTION = timedelta(days=7)
PURGE_INTERVAL_SECONDS = 3600 PURGE_INTERVAL_SECONDS = 3600
BATCH = 200
_SHARED_ELSEWHERE = """ _SHARED_ELSEWHERE = """
SELECT EXISTS (SELECT 1 FROM stories WHERE storage_key = $1) 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[]) 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( async def move_to_trash(
pool: asyncpg.Pool, pool: asyncpg.Pool,
@@ -86,12 +94,65 @@ async def move_to_trash(
return outcome 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( async def restore(
pool: asyncpg.Pool, pool: asyncpg.Pool,
storage: ContentAddressedStorage, storage: ContentAddressedStorage,
account_id: int, account_id: int,
keys: list[str], keys: list[str] | None,
) -> TrashOutcome: ) -> 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() outcome = TrashOutcome()
for key in keys: for key in keys:
async with pool.acquire() as conn, conn.transaction(): async with pool.acquire() as conn, conn.transaction():