feat(backend): storage trash with seven day retention, trashed files are not fetched again automatically

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)