FabledCurator/backend/app/services/backup_media.py

"""The off-site copy of the library's files: every image and attachment, sent
once each to the backup bucket as a content-addressed object.

Layout in the bucket, beside the database dumps under `database/`:

    media/ab/cd/<sha256><ext>

`ab` and `cd` are the first four hex digits of the file's sha256, which keeps
any one prefix small, and the extension keeps the files readable in the bucket
without FabledCurator. The sha256 is the one recorded at import, so a file
uploads exactly once however many rows or posts point at it, and a run is
incremental: its work list is every sha with no `backup_blob` row for the
current destination. A file removed from the library stays in the bucket until
the bucket's own retention removes it; FabledCurator never deletes there.

The bucket checks every object against that recorded sha256, so a file whose
bytes changed on disk since import is refused and reported, never stored under
the wrong name.
"""
from __future__ import annotations

import re
import time
from dataclasses import dataclass
from pathlib import Path

from sqlalchemy import delete, select, text
from sqlalchemy.dialects.postgresql import insert

from ..models import BackupBlob
from .backup_s3 import (
    Pacer,
    S3Error,
    S3Target,
    list_objects,
    object_present,
    put_file,
    put_file_in_parts,
)

PREFIX = "media/"
# Files above this size go in parts (S3 refuses a single PUT over 5 GiB, and a
# part that fails is resent alone instead of the whole file).
MULTIPART_OVER = 1 << 30
# The bucket is listed this many keys at a time during a reconcile, with a
# short rest between pages.
LIST_PAGE = 1000
LIST_REST_SECONDS = 0.2
_EXT = re.compile(r"^\.[a-z0-9]{1,10}$")


@dataclass(frozen=True)
class Blob:
    sha256: str
    path: str
    size_bytes: int
    mime: str | None


def destination_id(endpoint: str, bucket: str) -> str:
    """How backup_blob names a destination: one copy per endpoint and bucket."""
    return f"{endpoint.rstrip('/')}/{bucket}"


def target_id(target: S3Target) -> str:
    return destination_id(target.endpoint, target.bucket)


def blob_key(sha256: str, path: str) -> str:
    """The bucket key for a file: `media/ab/cd/<sha256><ext>`. An extension
    that isn't short and plain is left off, so a key is always safe."""
    ext = Path(path).suffix.lower()
    return f"{PREFIX}{sha256[:2]}/{sha256[2:4]}/{sha256}{ext if _EXT.match(ext) else ''}"


# Every library file once per sha256: images and attachments together. An
# attachment's sha can repeat across posts, and an image extracted from an
# archive is its own file with its own sha.
_LIBRARY = """
    SELECT sha256, path, size_bytes, mime FROM image_record
    UNION ALL
    SELECT sha256, path, size_bytes, mime FROM post_attachment
"""


def pending(session, tid: str, *, after: str = "", limit: int = 500) -> list[Blob]:
    """The next `limit` files not yet in the bucket, in sha order after
    `after`, so a run walks the library once even when some files fail."""
    rows = session.execute(text(f"""
        SELECT DISTINCT ON (lib.sha256) lib.sha256, lib.path, lib.size_bytes, lib.mime
        FROM ({_LIBRARY}) AS lib
        WHERE lib.sha256 > :after
          AND NOT EXISTS (
              SELECT 1 FROM backup_blob b WHERE b.target = :tid AND b.sha256 = lib.sha256
          )
        ORDER BY lib.sha256, lib.path
        LIMIT :limit
    """), {"after": after, "tid": tid, "limit": limit}).all()
    return [Blob(r.sha256, r.path, r.size_bytes, r.mime) for r in rows]


def coverage(session, tid: str | None) -> dict:
    """How much of the library is in the bucket: files and bytes, sent and
    total, each sha counted once."""
    row = session.execute(text(f"""
        WITH lib AS (
            SELECT DISTINCT ON (sha256) sha256, size_bytes FROM ({_LIBRARY}) AS u
            ORDER BY sha256
        )
        SELECT count(*) AS files_total,
               coalesce(sum(lib.size_bytes), 0) AS bytes_total,
               count(b.sha256) AS files_sent,
               coalesce(sum(lib.size_bytes) FILTER (WHERE b.sha256 IS NOT NULL), 0)
                   AS bytes_sent
        FROM lib
        LEFT JOIN backup_blob b ON b.sha256 = lib.sha256 AND b.target = :tid
    """), {"tid": tid or ""}).one()
    return {
        "files_total": row.files_total, "bytes_total": int(row.bytes_total),
        "files_sent": row.files_sent, "bytes_sent": int(row.bytes_sent),
    }


def send(target: S3Target, blob: Blob, images_root: Path, pacer: Pacer | None) -> str:
    """Upload one file to its key and return the key. Raises S3Error (or
    OSError) with the reason when it can't be sent."""
    path = Path(blob.path)
    if not path.is_absolute():
        path = images_root / path
    if not path.is_file():
        raise S3Error(f"{path} is not on disk.")
    size = path.stat().st_size
    if size != blob.size_bytes:
        raise S3Error(f"{path} is {size} bytes on disk but {blob.size_bytes} when imported.")
    key = blob_key(blob.sha256, blob.path)
    content_type = blob.mime or "application/octet-stream"
    try:
        if size > MULTIPART_OVER:
            put_file_in_parts(target, key, path, sha256=blob.sha256,
                              content_type=content_type, pacer=pacer)
        else:
            put_file(target, key, path, content_type=content_type,
                     payload_sha256=blob.sha256, pacer=pacer)
    except S3Error:
        if object_present(target, key, size):
            return key          # already there from an earlier run
        raise
    return key


def record(session, tid: str, blobs: list[tuple[Blob, str]]) -> None:
    """Note files as sent: (blob, key) pairs. A row already there is kept."""
    if blobs:
        session.execute(insert(BackupBlob).values([
            {"target": tid, "sha256": b.sha256, "key": key, "size_bytes": b.size_bytes}
            for b, key in blobs
        ]).on_conflict_do_nothing())


def reconcile(session, target: S3Target) -> dict:
    """Check backup_blob against what the bucket actually holds under
    `media/`. A row whose object is missing or the wrong size is dropped, so
    the next run sends that file again; an object the bucket holds for a
    library file with no row (a restored database, a lost commit) is recorded,
    so it isn't sent twice. Returns the counts."""
    tid = target_id(target)
    held: dict[str, int] = {}
    token = None
    while True:
        page = list_objects(target, PREFIX, LIST_PAGE, continuation=token)
        held.update((o.key, o.size) for o in page.objects)
        token = page.next_token
        if not token:
            break
        time.sleep(LIST_REST_SECONDS)

    rows = session.execute(
        select(BackupBlob.sha256, BackupBlob.key, BackupBlob.size_bytes)
        .where(BackupBlob.target == tid)
    ).all()
    gone = [r.sha256 for r in rows if held.get(r.key) != r.size_bytes]
    for at in range(0, len(gone), 1000):
        session.execute(delete(BackupBlob).where(
            BackupBlob.target == tid, BackupBlob.sha256.in_(gone[at:at + 1000]),
        ))
    kept = {r.sha256 for r in rows} - set(gone)

    found = []
    for r in session.execute(text(f"""
        SELECT DISTINCT ON (sha256) sha256, path, size_bytes, mime
        FROM ({_LIBRARY}) AS lib ORDER BY sha256, path
    """)).all():
        if r.sha256 in kept:
            continue
        blob = Blob(r.sha256, r.path, r.size_bytes, r.mime)
        key = blob_key(blob.sha256, blob.path)
        if held.get(key) == blob.size_bytes:
            found.append((blob, key))
    for at in range(0, len(found), 1000):
        record(session, tid, found[at:at + 1000])
    session.commit()
    return {"listed": len(held), "dropped": len(gone), "adopted": len(found)}