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:
1 parent
945adcf879
commit
f646f29ff9
9 files changed
+389
-15
No files matched your search
@@ -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
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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,
|
||||||
|
)
|
||||||
@@ -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())
|
|
||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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():
|
||||||
|
|||||||
Reference in new issue
Block a user