From f2c245707a032dcbb348f17af69586b35d53284b Mon Sep 17 00:00:00 2001 From: h Date: Wed, 14 Oct 2026 02:49:53 +0200 Subject: [PATCH] feat(backend): storage trash with seven day retention, trashed files are not fetched again automatically --- .../versions/c9e2a7f4b1d8_media_trash.py | 57 +++++ backend/src/api/app.py | 9 + backend/src/api/routers/media.py | 4 +- backend/src/api/routers/storage.py | 107 +++++++++ .../src/userbot/modules/capture/repository.py | 44 +++- .../modules/jobs/handlers/fetch_media.py | 4 +- .../src/userbot/modules/media/downloader.py | 14 +- backend/src/utils/read/models.py | 49 +++++ backend/src/utils/read/storage.py | 204 ++++++++++++++++++ backend/src/utils/storage.py | 28 +++ backend/src/utils/trash.py | 166 ++++++++++++++ 11 files changed, 680 insertions(+), 6 deletions(-) create mode 100644 backend/migrations/versions/c9e2a7f4b1d8_media_trash.py create mode 100644 backend/src/api/routers/storage.py create mode 100644 backend/src/utils/read/storage.py create mode 100644 backend/src/utils/trash.py diff --git a/backend/migrations/versions/c9e2a7f4b1d8_media_trash.py b/backend/migrations/versions/c9e2a7f4b1d8_media_trash.py new file mode 100644 index 0000000..86c16bd --- /dev/null +++ b/backend/migrations/versions/c9e2a7f4b1d8_media_trash.py @@ -0,0 +1,57 @@ +"""media trash + +Revision ID: c9e2a7f4b1d8 +Revises: b6d1f4a82e57 +Create Date: 2026-10-14 03:00:00.000000 + +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "c9e2a7f4b1d8" +down_revision: str | None = "b6d1f4a82e57" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_TRASH = """ +CREATE TABLE media_trash ( + storage_key text PRIMARY KEY, + account_id integer NOT NULL, + chat_id bigint NOT NULL, + message_id bigint NOT NULL, + kind text NOT NULL, + mime text, + file_name text NOT NULL, + file_size bigint, + trashed_at timestamptz NOT NULL DEFAULT now() +) +""" + + +def upgrade() -> None: + op.execute( + "ALTER TABLE media ADD COLUMN evicted_at timestamptz, " + "ADD COLUMN evicted_key text" + ) + op.execute( + "CREATE INDEX ix_media_evicted_unique ON media (unique_id) " + "WHERE evicted_at IS NOT NULL" + ) + op.execute( + "CREATE INDEX ix_media_evicted_key ON media (evicted_key) " + "WHERE evicted_key IS NOT NULL" + ) + op.execute(_TRASH) + op.execute( + "CREATE INDEX ix_media_trash_account ON media_trash " + "(account_id, trashed_at DESC)" + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS media_trash") + op.execute("DROP INDEX IF EXISTS ix_media_evicted_key") + op.execute("DROP INDEX IF EXISTS ix_media_evicted_unique") + op.execute("ALTER TABLE media DROP COLUMN evicted_key, DROP COLUMN evicted_at") diff --git a/backend/src/api/app.py b/backend/src/api/app.py index 753c816..2ce8f32 100644 --- a/backend/src/api/app.py +++ b/backend/src/api/app.py @@ -1,3 +1,4 @@ +import asyncio from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from pathlib import Path @@ -34,12 +35,15 @@ from api.routers import ( search, shares, social, + storage, stories, watches, ) from dependencies.container import container from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS, NO_STORE_HEADERS from utils.env import env +from utils.storage import ContentAddressedStorage +from utils.trash import purge_loop if env.auth.token is None: msg = "AUTH__TOKEN is required for the api process" @@ -55,7 +59,11 @@ mcp_app = mcp.http_app(path="/") async def lifespan(app_: Starlette) -> AsyncGenerator[None]: pool = await container.get(asyncpg.Pool) await hub.start(pool) + purger = asyncio.create_task( + purge_loop(pool, await container.get(ContentAddressedStorage)) + ) yield + purger.cancel() await hub.stop() await app_.state.dishka_container.close() @@ -96,6 +104,7 @@ app.include_router(discover.router) app.include_router(annotations.router) app.include_router(watches.router) app.include_router(shares.router) +app.include_router(storage.router) app.include_router(files.router) app.mount("/mcp", mcp_app) diff --git a/backend/src/api/routers/media.py b/backend/src/api/routers/media.py index 38f8626..af230c9 100644 --- a/backend/src/api/routers/media.py +++ b/backend/src/api/routers/media.py @@ -54,7 +54,7 @@ async def serve_media_version( download: Download = False, ) -> FileResponse: version = await get_media_version(pool, version_id) - if version is None: + if version is None or not storage.exists(version.storage_key): raise HTTPException(status_code=404, detail="media version not found") headers = IMMUTABLE_HEADERS if download: @@ -96,6 +96,8 @@ async def serve_media( status_code=409, detail="media not downloaded; enqueue fetch via POST /api/media/fetch", ) + if not storage.exists(media.storage_key): + raise HTTPException(status_code=404, detail="media file is missing") headers = DAY_HEADERS if download: headers = _attachment(headers, media.file_name or f"media_{media_id}") diff --git a/backend/src/api/routers/storage.py b/backend/src/api/routers/storage.py new file mode 100644 index 0000000..92f9e27 --- /dev/null +++ b/backend/src/api/routers/storage.py @@ -0,0 +1,107 @@ +from typing import Annotated + +import asyncpg +from dishka.integrations.fastapi import DishkaRoute, FromDishka +from fastapi import APIRouter, HTTPException, Query +from fastapi.responses import FileResponse +from pydantic import BaseModel + +from utils import trash +from utils.cache import NO_STORE_HEADERS +from utils.files import resolve_mime +from utils.read import storage as storage_read +from utils.read.models import ( + DEFAULT_LIMIT, + Page, + StorageFileView, + StorageSummaryView, + TrashItemView, + TrashOutcome, +) +from utils.storage import ContentAddressedStorage + +router = APIRouter(prefix="/api/storage", tags=["storage"], route_class=DishkaRoute) + + +class TrashRequest(BaseModel): + account_id: int + media_ids: list[int] + + +class RestoreRequest(BaseModel): + account_id: int + keys: list[str] + + +class PurgeRequest(BaseModel): + account_id: int + keys: list[str] | None = None + + +@router.get("/summary") +async def summary( + pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()] +) -> StorageSummaryView: + return await storage_read.get_summary(pool, account_id) + + +@router.get("/files") +async def files( + 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[StorageFileView]: + page = Page(limit=limit, offset=offset) + return await storage_read.list_files(pool, account_id, group, page) + + +@router.get("/trash") +async def trash_items( + pool: FromDishka[asyncpg.Pool], account_id: Annotated[int, Query()] +) -> list[TrashItemView]: + return await storage_read.list_trash(pool, account_id) + + +@router.post("/trash") +async def move_to_trash( + pool: FromDishka[asyncpg.Pool], + storage: FromDishka[ContentAddressedStorage], + body: TrashRequest, +) -> TrashOutcome: + return await trash.move_to_trash(pool, storage, body.account_id, body.media_ids) + + +@router.post("/trash/restore") +async def restore( + pool: FromDishka[asyncpg.Pool], + storage: FromDishka[ContentAddressedStorage], + body: RestoreRequest, +) -> TrashOutcome: + return await trash.restore(pool, storage, body.account_id, body.keys) + + +@router.post("/trash/purge") +async def purge( + pool: FromDishka[asyncpg.Pool], + storage: FromDishka[ContentAddressedStorage], + body: PurgeRequest, +) -> TrashOutcome: + return await trash.purge(pool, storage, body.account_id, body.keys) + + +@router.get("/trash/{storage_key}/file") +async def serve_trashed( + pool: FromDishka[asyncpg.Pool], + storage: FromDishka[ContentAddressedStorage], + storage_key: str, +) -> FileResponse: + item = await storage_read.get_trash_item(pool, storage_key) + if item is None: + raise HTTPException(status_code=404, detail="file is not in trash") + return FileResponse( + storage.trash_url(storage_key), + media_type=resolve_mime(item["kind"], item["mime"], item["file_name"]), + headers=NO_STORE_HEADERS, + ) diff --git a/backend/src/userbot/modules/capture/repository.py b/backend/src/userbot/modules/capture/repository.py index b137e32..e87c327 100644 --- a/backend/src/userbot/modules/capture/repository.py +++ b/backend/src/userbot/modules/capture/repository.py @@ -88,8 +88,8 @@ ON CONFLICT (account_id, chat_id, message_id, storage_key) DO NOTHING _INSERT_MEDIA = """ INSERT INTO media (account_id, chat_id, message_id, kind, storage_key, file_size, mime, - ttl_seconds, downloaded, unique_id) -VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + ttl_seconds, downloaded, unique_id, evicted_at) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, CASE WHEN $11 THEN now() END) ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET kind = EXCLUDED.kind, storage_key = EXCLUDED.storage_key, @@ -97,7 +97,27 @@ ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET mime = EXCLUDED.mime, ttl_seconds = EXCLUDED.ttl_seconds, downloaded = EXCLUDED.downloaded, - unique_id = EXCLUDED.unique_id + unique_id = EXCLUDED.unique_id, + evicted_at = CASE + WHEN EXCLUDED.downloaded THEN NULL + WHEN media.unique_id IS DISTINCT FROM EXCLUDED.unique_id + THEN EXCLUDED.evicted_at + ELSE coalesce(media.evicted_at, EXCLUDED.evicted_at) END, + evicted_key = CASE WHEN EXCLUDED.downloaded + OR media.unique_id IS DISTINCT FROM EXCLUDED.unique_id + THEN NULL ELSE media.evicted_key END +""" + +_IS_EVICTED = """ +SELECT EXISTS ( + SELECT 1 FROM media WHERE unique_id = $1 AND evicted_at IS NOT NULL +) +""" + +_UNEVICT = """ +UPDATE media SET storage_key = $2, file_size = $3, downloaded = true, + evicted_at = NULL, evicted_key = NULL +WHERE evicted_key = $2 OR (unique_id = $1 AND evicted_at IS NOT NULL) """ @@ -229,6 +249,7 @@ async def insert_media( # noqa: PLR0913 unique_id: str | None, *, downloaded: bool, + evicted: bool = False, ) -> None: async with pool.acquire() as conn, conn.transaction(): await conn.execute( @@ -243,6 +264,7 @@ async def insert_media( # noqa: PLR0913 ttl_seconds, downloaded, unique_id, + evicted, ) if storage_key is not None: await conn.execute( @@ -259,6 +281,22 @@ async def insert_media( # noqa: PLR0913 ) +async def is_evicted(pool: asyncpg.Pool, unique_id: str | None) -> bool: + return unique_id is not None and await pool.fetchval(_IS_EVICTED, unique_id) + + +async def unevict( + pool: asyncpg.Pool, unique_id: str | None, storage_key: str, file_size: int | None +) -> bool: + async with pool.acquire() as conn, conn.transaction(): + await conn.execute(_UNEVICT, unique_id, storage_key, file_size) + dropped = await conn.fetchval( + "DELETE FROM media_trash WHERE storage_key = $1 RETURNING storage_key", + storage_key, + ) + return dropped is not None + + async def insert_callbacks( pool: asyncpg.Pool, account_id: int, diff --git a/backend/src/userbot/modules/jobs/handlers/fetch_media.py b/backend/src/userbot/modules/jobs/handlers/fetch_media.py index 6bb5ef3..7f6e7bf 100644 --- a/backend/src/userbot/modules/jobs/handlers/fetch_media.py +++ b/backend/src/userbot/modules/jobs/handlers/fetch_media.py @@ -21,7 +21,9 @@ async def fetch_media(ctx: JobContext) -> None: msg = f"message {chat_id}/{message_id} is unavailable" raise RuntimeError(msg) toggles = CaptureToggles(media=True, self_destruct_media=True) - ok = await capture_media(client, message, capture, chat_id, message_id, toggles) + ok = await capture_media( + client, message, capture, chat_id, message_id, toggles, force=True + ) if not ok: msg = f"media download failed for {chat_id}/{message_id}" raise RuntimeError(msg) diff --git a/backend/src/userbot/modules/media/downloader.py b/backend/src/userbot/modules/media/downloader.py index d488fe1..5a1d2d8 100644 --- a/backend/src/userbot/modules/media/downloader.py +++ b/backend/src/userbot/modules/media/downloader.py @@ -106,6 +106,8 @@ async def capture_media( # noqa: PLR0913 chat_id: int, message_id: int, toggles: CaptureToggles, + *, + force: bool = False, ) -> bool: kind, obj = media_object(message) if obj is None: @@ -117,6 +119,7 @@ async def capture_media( # noqa: PLR0913 mime = getattr(obj, "mime_type", None) storage_key: str | None = None downloaded = False + evicted = False flood: FloodWait | FloodPremiumWait | None = None if want: existing = await repository.current_media( @@ -131,7 +134,7 @@ async def capture_media( # noqa: PLR0913 storage_key = existing["storage_key"] file_size = existing["file_size"] downloaded = True - else: + elif force or not await repository.is_evicted(ctx.pool, unique_id): target = _download_target(message) try: data = await _download(client, target, chat_id, message_id) @@ -142,6 +145,8 @@ async def capture_media( # noqa: PLR0913 storage_key = ctx.storage.put(data) file_size = len(data) downloaded = True + else: + evicted = True await repository.insert_media( ctx.pool, ctx.account_id, @@ -154,7 +159,14 @@ async def capture_media( # noqa: PLR0913 ttl, unique_id, downloaded=downloaded, + evicted=evicted, ) + if ( + force + and storage_key is not None + and await repository.unevict(ctx.pool, unique_id, storage_key, file_size) + ): + ctx.storage.purge(storage_key) if flood is not None: raise flood return downloaded or not want diff --git a/backend/src/utils/read/models.py b/backend/src/utils/read/models.py index f907416..d409080 100644 --- a/backend/src/utils/read/models.py +++ b/backend/src/utils/read/models.py @@ -236,6 +236,55 @@ class MediaVersionView(BaseModel): observed_at: datetime +class StorageGroupView(BaseModel): + group: str + files: int + bytes: int + + +class StorageSummaryView(BaseModel): + files: int + bytes: int + heavy_files: int + heavy_bytes: int + heavy_threshold: int + trash_files: int + trash_bytes: int + groups: list[StorageGroupView] + + +class StorageFileView(BaseModel): + media_id: int + chat_id: int + message_id: int + kind: str + mime: str | None + file_name: str + file_size: int + date: datetime | None + refs: int + refetchable: bool + + +class TrashItemView(BaseModel): + storage_key: str + chat_id: int + message_id: int + kind: str + mime: str | None + file_name: str + file_size: int | None + trashed_at: datetime + purge_at: datetime + refetchable: bool + + +class TrashOutcome(BaseModel): + files: int = 0 + bytes: int = 0 + blocked: list[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 new file mode 100644 index 0000000..a43e994 --- /dev/null +++ b/backend/src/utils/read/storage.py @@ -0,0 +1,204 @@ +import asyncio +from collections import defaultdict + +import asyncpg + +from utils.files import media_file_name +from utils.read.models import ( + Page, + StorageFileView, + StorageGroupView, + StorageSummaryView, + TrashItemView, +) +from utils.trash import RETENTION + +HEAVY_BYTES = 100 * 1024 * 1024 +REFS_CHECKED = 5 +LOOKUP_CONCURRENCY = 8 + +OTHER_GROUP = "other" +GROUP_KINDS: dict[str, tuple[str, ...]] = { + "video": ("video",), + "file": ("document",), + "voice": ("video_note", "voice"), + "audio": ("audio",), +} +_KIND_GROUP = {kind: group for group, kinds in GROUP_KINDS.items() for kind in kinds} + +_SUMMARY = """ +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 + GROUP BY storage_key +) +SELECT kind, count(*) AS files, coalesce(sum(size), 0)::bigint AS bytes, + count(*) FILTER (WHERE size >= $2) AS heavy_files, + coalesce(sum(size) FILTER (WHERE size >= $2), 0)::bigint AS heavy_bytes +FROM blobs GROUP BY kind +""" + +_TRASH_TOTAL = """ +SELECT count(*) AS files, coalesce(sum(file_size), 0)::bigint AS bytes +FROM media_trash WHERE account_id = $1 +""" + +_FILES = """ +WITH top AS MATERIALIZED ( + SELECT storage_key, max(file_size) AS file_size, count(*) AS refs + 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 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, + m.kind, m.mime, m.ttl_seconds, + row_number() OVER (PARTITION BY t.storage_key ORDER BY 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 +""" + +_TRASH = """ +SELECT storage_key, chat_id, message_id, kind, mime, file_name, file_size, + trashed_at +FROM media_trash WHERE account_id = $1 +ORDER BY trashed_at DESC, storage_key +""" + +_TRASH_REFS = """ +SELECT * FROM ( + SELECT evicted_key AS storage_key, chat_id, message_id, kind, ttl_seconds, + row_number() OVER (PARTITION BY evicted_key ORDER BY id) AS n + FROM media + WHERE account_id = $1 AND evicted_key = ANY($2::text[]) +) ranked +WHERE n <= $3 +""" + +_MESSAGE = """ +SELECT date, deleted_at, raw->$4::text->>'file_name' AS original_name +FROM messages +WHERE account_id = $1 AND chat_id = $2 AND message_id = $3 +ORDER BY date DESC LIMIT 1 +""" + + +def _group_filter(group: str | None) -> tuple[list[str] | None, bool]: + if group is None: + return None, True + if group == OTHER_GROUP: + return list(_KIND_GROUP), False + 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) + groups = { + name: StorageGroupView(group=name, files=0, bytes=0) + for name in (*GROUP_KINDS, OTHER_GROUP) + } + for row in rows: + group = groups[_KIND_GROUP.get(row["kind"], OTHER_GROUP)] + group.files += row["files"] + group.bytes += row["bytes"] + return StorageSummaryView( + files=sum(row["files"] for row in rows), + bytes=sum(row["bytes"] for row in rows), + heavy_files=sum(row["heavy_files"] for row in rows), + heavy_bytes=sum(row["heavy_bytes"] for row in rows), + heavy_threshold=HEAVY_BYTES, + trash_files=trash["files"], + trash_bytes=trash["bytes"], + groups=list(groups.values()), + ) + + +async def _messages( + pool: asyncpg.Pool, account_id: int, refs: list[asyncpg.Record] +) -> list[asyncpg.Record | None]: + gate = asyncio.Semaphore(LOOKUP_CONCURRENCY) + + async def lookup(ref: asyncpg.Record) -> asyncpg.Record | None: + async with gate: + return await pool.fetchrow( + _MESSAGE, account_id, ref["chat_id"], ref["message_id"], ref["kind"] + ) + + return await asyncio.gather(*(lookup(ref) for ref in refs)) + + +def _alive(ref: asyncpg.Record, message: asyncpg.Record | None) -> bool: + return ( + message is not None + and message["deleted_at"] is None + and ref["ttl_seconds"] is None + ) + + +async def list_files( + pool: asyncpg.Pool, account_id: int, group: str | None, page: Page +) -> list[StorageFileView]: + kinds, include = _group_filter(group) + refs = await pool.fetch( + _FILES, account_id, kinds, include, page.capped_limit, page.offset, REFS_CHECKED + ) + 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): + continue + files[ref["storage_key"]] = StorageFileView( + media_id=ref["id"], + chat_id=ref["chat_id"], + message_id=ref["message_id"], + kind=ref["kind"], + mime=ref["mime"], + file_name=media_file_name( + ref["kind"], + ref["mime"], + ref["message_id"], + message["original_name"] if message else None, + ), + file_size=ref["file_size"] or 0, + date=message["date"] if message else None, + refs=ref["refs"], + refetchable=alive, + ) + return list(files.values()) + + +async def list_trash(pool: asyncpg.Pool, account_id: int) -> list[TrashItemView]: + rows = await pool.fetch(_TRASH, account_id) + 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) + for ref, message in zip(refs, messages, strict=True): + alive[ref["storage_key"]] |= _alive(ref, message) + return [ + TrashItemView( + **dict(row), + purge_at=row["trashed_at"] + RETENTION, + refetchable=alive.get(row["storage_key"], True), + ) + for row in rows + ] + + +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", + storage_key, + ) diff --git a/backend/src/utils/storage.py b/backend/src/utils/storage.py index f7e7786..b34b4d5 100644 --- a/backend/src/utils/storage.py +++ b/backend/src/utils/storage.py @@ -1,6 +1,8 @@ import hashlib from pathlib import Path +TRASH_DIR = ".trash" + class ContentAddressedStorage: def __init__(self, root: str | Path, shard_depth: int = 2) -> None: @@ -27,3 +29,29 @@ class ContentAddressedStorage: def url(self, key: str) -> str: return str(self._path(key)) + + def _trash_path(self, key: str) -> Path: + return self._root / TRASH_DIR / key + + def trash(self, key: str) -> None: + target = self._trash_path(key) + target.parent.mkdir(parents=True, exist_ok=True) + self._path(key).replace(target) + + def restore(self, key: str) -> bool: + source = self._trash_path(key) + target = self._path(key) + if not source.exists(): + return target.exists() + if target.exists(): + source.unlink() + else: + target.parent.mkdir(parents=True, exist_ok=True) + source.replace(target) + return True + + def purge(self, key: str) -> None: + self._trash_path(key).unlink(missing_ok=True) + + def trash_url(self, key: str) -> str: + return str(self._trash_path(key)) diff --git a/backend/src/utils/trash.py b/backend/src/utils/trash.py new file mode 100644 index 0000000..0ed6e57 --- /dev/null +++ b/backend/src/utils/trash.py @@ -0,0 +1,166 @@ +import asyncio +import logging +from datetime import timedelta + +import asyncpg + +from utils.read.media import get_media +from utils.read.models import TrashOutcome +from utils.storage import ContentAddressedStorage + +logger = logging.getLogger(__name__) + +RETENTION = timedelta(days=7) +PURGE_INTERVAL_SECONDS = 3600 + +_SHARED_ELSEWHERE = """ +SELECT EXISTS (SELECT 1 FROM stories WHERE storage_key = $1) + OR EXISTS (SELECT 1 FROM avatars WHERE storage_key = $1) + OR EXISTS (SELECT 1 FROM custom_emoji WHERE storage_key = $1) +""" + +_EVICT = """ +UPDATE media SET storage_key = NULL, downloaded = false, + evicted_at = now(), evicted_key = $2 +WHERE storage_key = $1 +""" + +_INSERT = """ +INSERT INTO media_trash + (storage_key, account_id, chat_id, message_id, kind, mime, file_name, file_size) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8) +ON CONFLICT (storage_key) DO UPDATE SET trashed_at = now() +""" + +_RESTORE = """ +UPDATE media SET storage_key = $1, downloaded = true, + evicted_at = NULL, evicted_key = NULL +WHERE evicted_key = $1 OR ( + evicted_at IS NOT NULL AND unique_id IN ( + SELECT unique_id FROM media + WHERE evicted_key = $1 AND unique_id IS NOT NULL + ) +) +""" + +_FORGET_KEYS = """ +UPDATE media SET evicted_key = NULL WHERE evicted_key = ANY($1::text[]) +""" + + +async def move_to_trash( + pool: asyncpg.Pool, + storage: ContentAddressedStorage, + account_id: int, + media_ids: list[int], +) -> TrashOutcome: + outcome = TrashOutcome() + for media_id in media_ids: + media = await get_media(pool, media_id) + if media is None or media.account_id != account_id or not media.storage_key: + continue + key = media.storage_key + on_disk = storage.exists(key) + async with pool.acquire() as conn, conn.transaction(): + if await conn.fetchval(_SHARED_ELSEWHERE, key): + outcome.blocked.append(media_id) + continue + await conn.execute(_EVICT, key, key if on_disk else None) + if not on_disk: + continue + await conn.execute( + _INSERT, + key, + account_id, + media.chat_id, + media.message_id, + media.kind, + media.mime, + media.file_name, + media.file_size, + ) + storage.trash(key) + outcome.files += 1 + outcome.bytes += media.file_size or 0 + return outcome + + +async def restore( + pool: asyncpg.Pool, + storage: ContentAddressedStorage, + account_id: int, + keys: list[str], +) -> TrashOutcome: + outcome = TrashOutcome() + for key in keys: + async with pool.acquire() as conn, conn.transaction(): + row = await conn.fetchrow( + "DELETE FROM media_trash WHERE storage_key = $1 AND account_id = $2 " + "RETURNING file_size", + key, + account_id, + ) + if row is None: + continue + if not storage.restore(key): + await conn.execute(_FORGET_KEYS, [key]) + continue + await conn.execute(_RESTORE, key) + outcome.files += 1 + outcome.bytes += row["file_size"] or 0 + return outcome + + +async def _purge( + pool: asyncpg.Pool, storage: ContentAddressedStorage, rows: list[asyncpg.Record] +) -> TrashOutcome: + keys = [row["storage_key"] for row in rows] + for key in keys: + storage.purge(key) + async with pool.acquire() as conn, conn.transaction(): + await conn.execute( + "DELETE FROM media_trash WHERE storage_key = ANY($1::text[])", keys + ) + await conn.execute(_FORGET_KEYS, keys) + return TrashOutcome( + files=len(rows), bytes=sum(row["file_size"] or 0 for row in rows) + ) + + +async def purge( + pool: asyncpg.Pool, + storage: ContentAddressedStorage, + account_id: int, + keys: list[str] | None, +) -> TrashOutcome: + rows = await pool.fetch( + "SELECT storage_key, file_size FROM media_trash WHERE account_id = $1 " + "AND ($2::text[] IS NULL OR storage_key = ANY($2))", + account_id, + keys, + ) + return await _purge(pool, storage, rows) + + +async def purge_expired( + pool: asyncpg.Pool, storage: ContentAddressedStorage +) -> TrashOutcome: + rows = await pool.fetch( + "SELECT storage_key, file_size FROM media_trash " + "WHERE trashed_at < now() - $1::interval", + RETENTION, + ) + return await _purge(pool, storage, rows) + + +async def purge_loop(pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None: + while True: + try: + outcome = await purge_expired(pool, storage) + if outcome.files: + logger.info( + "Trash purged: %s files, %s bytes", outcome.files, outcome.bytes + ) + except Exception: + logger.exception("Trash purge failed") + await asyncio.sleep(PURGE_INTERVAL_SECONDS)