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