feat(backend): storage trash with seven day retention, trashed files are not fetched again automatically
This commit is contained in:
hh committed
1 parent
c20e132872
commit
f2c245707a
11 files changed
+680
-6
No files matched your search
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user