229 lines
6.9 KiB
Python
229 lines
6.9 KiB
Python
import asyncio
|
|
import logging
|
|
from datetime import timedelta
|
|
|
|
import asyncpg
|
|
|
|
from utils.files import media_file_name
|
|
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
|
|
BATCH = 200
|
|
|
|
_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[])
|
|
"""
|
|
|
|
_EVICT_MANY = """
|
|
UPDATE media SET evicted_key = storage_key, storage_key = NULL,
|
|
downloaded = false, evicted_at = now()
|
|
WHERE storage_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.keys.append(key)
|
|
outcome.files += 1
|
|
outcome.bytes += media.file_size or 0
|
|
return outcome
|
|
|
|
|
|
def _move_batch(storage: ContentAddressedStorage, keys: list[str]) -> None:
|
|
moved: list[str] = []
|
|
try:
|
|
for key in keys:
|
|
if storage.exists(key):
|
|
storage.trash(key)
|
|
moved.append(key)
|
|
except OSError:
|
|
for key in moved:
|
|
storage.restore(key)
|
|
raise
|
|
|
|
|
|
async def trash_blobs(
|
|
pool: asyncpg.Pool,
|
|
storage: ContentAddressedStorage,
|
|
account_id: int,
|
|
blobs: list[asyncpg.Record],
|
|
) -> TrashOutcome:
|
|
outcome = TrashOutcome()
|
|
for start in range(0, len(blobs), BATCH):
|
|
batch = blobs[start : start + BATCH]
|
|
keys = [blob["storage_key"] for blob in batch]
|
|
async with pool.acquire() as conn, conn.transaction():
|
|
await conn.execute(_EVICT_MANY, keys)
|
|
await conn.executemany(
|
|
_INSERT,
|
|
[
|
|
(
|
|
blob["storage_key"],
|
|
account_id,
|
|
blob["chat_id"],
|
|
blob["message_id"],
|
|
blob["kind"],
|
|
blob["mime"],
|
|
media_file_name(blob["kind"], blob["mime"], blob["message_id"]),
|
|
blob["file_size"],
|
|
)
|
|
for blob in batch
|
|
],
|
|
)
|
|
await asyncio.to_thread(_move_batch, storage, keys)
|
|
outcome.keys.extend(keys)
|
|
outcome.files += len(batch)
|
|
outcome.bytes += sum(blob["file_size"] or 0 for blob in batch)
|
|
return outcome
|
|
|
|
|
|
async def restore(
|
|
pool: asyncpg.Pool,
|
|
storage: ContentAddressedStorage,
|
|
account_id: int,
|
|
keys: list[str] | None,
|
|
) -> TrashOutcome:
|
|
if keys is None:
|
|
rows = await pool.fetch(
|
|
"SELECT storage_key FROM media_trash WHERE account_id = $1", account_id
|
|
)
|
|
keys = [row["storage_key"] for row in rows]
|
|
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)
|