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 collections.abc import AsyncGenerator
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -34,12 +35,15 @@ from api.routers import (
|
|||||||
search,
|
search,
|
||||||
shares,
|
shares,
|
||||||
social,
|
social,
|
||||||
|
storage,
|
||||||
stories,
|
stories,
|
||||||
watches,
|
watches,
|
||||||
)
|
)
|
||||||
from dependencies.container import container
|
from dependencies.container import container
|
||||||
from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS, NO_STORE_HEADERS
|
from utils.cache import DAY_HEADERS, IMMUTABLE_HEADERS, NO_STORE_HEADERS
|
||||||
from utils.env import env
|
from utils.env import env
|
||||||
|
from utils.storage import ContentAddressedStorage
|
||||||
|
from utils.trash import purge_loop
|
||||||
|
|
||||||
if env.auth.token is None:
|
if env.auth.token is None:
|
||||||
msg = "AUTH__TOKEN is required for the api process"
|
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]:
|
async def lifespan(app_: Starlette) -> AsyncGenerator[None]:
|
||||||
pool = await container.get(asyncpg.Pool)
|
pool = await container.get(asyncpg.Pool)
|
||||||
await hub.start(pool)
|
await hub.start(pool)
|
||||||
|
purger = asyncio.create_task(
|
||||||
|
purge_loop(pool, await container.get(ContentAddressedStorage))
|
||||||
|
)
|
||||||
yield
|
yield
|
||||||
|
purger.cancel()
|
||||||
await hub.stop()
|
await hub.stop()
|
||||||
await app_.state.dishka_container.close()
|
await app_.state.dishka_container.close()
|
||||||
|
|
||||||
@@ -96,6 +104,7 @@ app.include_router(discover.router)
|
|||||||
app.include_router(annotations.router)
|
app.include_router(annotations.router)
|
||||||
app.include_router(watches.router)
|
app.include_router(watches.router)
|
||||||
app.include_router(shares.router)
|
app.include_router(shares.router)
|
||||||
|
app.include_router(storage.router)
|
||||||
app.include_router(files.router)
|
app.include_router(files.router)
|
||||||
|
|
||||||
app.mount("/mcp", mcp_app)
|
app.mount("/mcp", mcp_app)
|
||||||
|
|||||||
@@ -54,7 +54,7 @@ async def serve_media_version(
|
|||||||
download: Download = False,
|
download: Download = False,
|
||||||
) -> FileResponse:
|
) -> FileResponse:
|
||||||
version = await get_media_version(pool, version_id)
|
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")
|
raise HTTPException(status_code=404, detail="media version not found")
|
||||||
headers = IMMUTABLE_HEADERS
|
headers = IMMUTABLE_HEADERS
|
||||||
if download:
|
if download:
|
||||||
@@ -96,6 +96,8 @@ async def serve_media(
|
|||||||
status_code=409,
|
status_code=409,
|
||||||
detail="media not downloaded; enqueue fetch via POST /api/media/fetch",
|
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
|
headers = DAY_HEADERS
|
||||||
if download:
|
if download:
|
||||||
headers = _attachment(headers, media.file_name or f"media_{media_id}")
|
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_MEDIA = """
|
||||||
INSERT INTO media
|
INSERT INTO media
|
||||||
(account_id, chat_id, message_id, kind, storage_key, file_size, mime,
|
(account_id, chat_id, message_id, kind, storage_key, file_size, mime,
|
||||||
ttl_seconds, downloaded, unique_id)
|
ttl_seconds, downloaded, unique_id, evicted_at)
|
||||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
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
|
ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
|
||||||
kind = EXCLUDED.kind,
|
kind = EXCLUDED.kind,
|
||||||
storage_key = EXCLUDED.storage_key,
|
storage_key = EXCLUDED.storage_key,
|
||||||
@@ -97,7 +97,27 @@ ON CONFLICT (account_id, chat_id, message_id) DO UPDATE SET
|
|||||||
mime = EXCLUDED.mime,
|
mime = EXCLUDED.mime,
|
||||||
ttl_seconds = EXCLUDED.ttl_seconds,
|
ttl_seconds = EXCLUDED.ttl_seconds,
|
||||||
downloaded = EXCLUDED.downloaded,
|
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,
|
unique_id: str | None,
|
||||||
*,
|
*,
|
||||||
downloaded: bool,
|
downloaded: bool,
|
||||||
|
evicted: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
async with pool.acquire() as conn, conn.transaction():
|
async with pool.acquire() as conn, conn.transaction():
|
||||||
await conn.execute(
|
await conn.execute(
|
||||||
@@ -243,6 +264,7 @@ async def insert_media( # noqa: PLR0913
|
|||||||
ttl_seconds,
|
ttl_seconds,
|
||||||
downloaded,
|
downloaded,
|
||||||
unique_id,
|
unique_id,
|
||||||
|
evicted,
|
||||||
)
|
)
|
||||||
if storage_key is not None:
|
if storage_key is not None:
|
||||||
await conn.execute(
|
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(
|
async def insert_callbacks(
|
||||||
pool: asyncpg.Pool,
|
pool: asyncpg.Pool,
|
||||||
account_id: int,
|
account_id: int,
|
||||||
|
|||||||
@@ -21,7 +21,9 @@ async def fetch_media(ctx: JobContext) -> None:
|
|||||||
msg = f"message {chat_id}/{message_id} is unavailable"
|
msg = f"message {chat_id}/{message_id} is unavailable"
|
||||||
raise RuntimeError(msg)
|
raise RuntimeError(msg)
|
||||||
toggles = CaptureToggles(media=True, self_destruct_media=True)
|
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:
|
if not ok:
|
||||||
msg = f"media download failed for {chat_id}/{message_id}"
|
msg = f"media download failed for {chat_id}/{message_id}"
|
||||||
raise RuntimeError(msg)
|
raise RuntimeError(msg)
|
||||||
@@ -106,6 +106,8 @@ async def capture_media( # noqa: PLR0913
|
|||||||
chat_id: int,
|
chat_id: int,
|
||||||
message_id: int,
|
message_id: int,
|
||||||
toggles: CaptureToggles,
|
toggles: CaptureToggles,
|
||||||
|
*,
|
||||||
|
force: bool = False,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
kind, obj = media_object(message)
|
kind, obj = media_object(message)
|
||||||
if obj is None:
|
if obj is None:
|
||||||
@@ -117,6 +119,7 @@ async def capture_media( # noqa: PLR0913
|
|||||||
mime = getattr(obj, "mime_type", None)
|
mime = getattr(obj, "mime_type", None)
|
||||||
storage_key: str | None = None
|
storage_key: str | None = None
|
||||||
downloaded = False
|
downloaded = False
|
||||||
|
evicted = False
|
||||||
flood: FloodWait | FloodPremiumWait | None = None
|
flood: FloodWait | FloodPremiumWait | None = None
|
||||||
if want:
|
if want:
|
||||||
existing = await repository.current_media(
|
existing = await repository.current_media(
|
||||||
@@ -131,7 +134,7 @@ async def capture_media( # noqa: PLR0913
|
|||||||
storage_key = existing["storage_key"]
|
storage_key = existing["storage_key"]
|
||||||
file_size = existing["file_size"]
|
file_size = existing["file_size"]
|
||||||
downloaded = True
|
downloaded = True
|
||||||
else:
|
elif force or not await repository.is_evicted(ctx.pool, unique_id):
|
||||||
target = _download_target(message)
|
target = _download_target(message)
|
||||||
try:
|
try:
|
||||||
data = await _download(client, target, chat_id, message_id)
|
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)
|
storage_key = ctx.storage.put(data)
|
||||||
file_size = len(data)
|
file_size = len(data)
|
||||||
downloaded = True
|
downloaded = True
|
||||||
|
else:
|
||||||
|
evicted = True
|
||||||
await repository.insert_media(
|
await repository.insert_media(
|
||||||
ctx.pool,
|
ctx.pool,
|
||||||
ctx.account_id,
|
ctx.account_id,
|
||||||
@@ -154,7 +159,14 @@ async def capture_media( # noqa: PLR0913
|
|||||||
ttl,
|
ttl,
|
||||||
unique_id,
|
unique_id,
|
||||||
downloaded=downloaded,
|
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:
|
if flood is not None:
|
||||||
raise flood
|
raise flood
|
||||||
return downloaded or not want
|
return downloaded or not want
|
||||||
@@ -236,6 +236,55 @@ class MediaVersionView(BaseModel):
|
|||||||
observed_at: datetime
|
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):
|
class AvatarRef(BaseModel):
|
||||||
unique_id: str
|
unique_id: str
|
||||||
storage_key: str | None
|
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
|
import hashlib
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
TRASH_DIR = ".trash"
|
||||||
|
|
||||||
|
|
||||||
class ContentAddressedStorage:
|
class ContentAddressedStorage:
|
||||||
def __init__(self, root: str | Path, shard_depth: int = 2) -> None:
|
def __init__(self, root: str | Path, shard_depth: int = 2) -> None:
|
||||||
@@ -27,3 +29,29 @@ class ContentAddressedStorage:
|
|||||||
|
|
||||||
def url(self, key: str) -> str:
|
def url(self, key: str) -> str:
|
||||||
return str(self._path(key))
|
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