api/storage.py

"""Cloudflare R2 (S3-compatible) access: uploads from the worker, presigned links for the report page."""

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

CONTENT_TYPES = {
    ".mcap": "application/octet-stream",
    ".mp4": "video/mp4",
    ".json": "application/json",
    ".xml": "application/xml",
    ".webm": "video/webm",
    ".npz": "application/octet-stream",
    ".log": "text/plain; charset=utf-8",
}


class Storage(Protocol):
    def upload(self, local: Path, key: str) -> None: ...

    def url(self, key: str, expires_s: int = 7 * 24 * 3600) -> str: ...

    def read_text(self, key: str) -> str | None: ...

    def total_bytes(self) -> int: ...


@dataclass(frozen=True)
class R2Config:
    account_id: str
    access_key_id: str
    secret_access_key: str
    bucket: str

    @classmethod
    def from_env(cls) -> R2Config | None:
        values = [
            os.environ.get(k, "") for k in ("R2_ACCOUNT_ID", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY", "R2_BUCKET")
        ]
        return cls(*values) if all(values) else None


class R2Storage:
    def __init__(self, config: R2Config):
        import boto3
        from botocore.config import Config

        self.bucket = config.bucket
        self._s3 = boto3.client(
            "s3",
            endpoint_url=f"https://{config.account_id}.r2.cloudflarestorage.com",
            aws_access_key_id=config.access_key_id,
            aws_secret_access_key=config.secret_access_key,
            region_name="auto",
            config=Config(signature_version="s3v4", retries={"max_attempts": 5}),
        )

    def upload(self, local: Path, key: str) -> None:
        extra = {"ContentType": CONTENT_TYPES.get(Path(key).suffix, "application/octet-stream")}
        self._s3.upload_file(str(local), self.bucket, key, ExtraArgs=extra)

    def url(self, key: str, expires_s: int = 7 * 24 * 3600) -> str:
        # S3 SigV4 presigned URLs are valid for at most 7 days.
        return self._s3.generate_presigned_url(
            "get_object", Params={"Bucket": self.bucket, "Key": key}, ExpiresIn=min(expires_s, 7 * 24 * 3600)
        )

    def read_text(self, key: str) -> str | None:
        try:
            return self._s3.get_object(Bucket=self.bucket, Key=key)["Body"].read().decode()
        except self._s3.exceptions.NoSuchKey:
            return None

    def total_bytes(self) -> int:
        """Everything stored in the bucket (one listing request per 1000 files)."""
        total = 0
        for page in self._s3.get_paginator("list_objects_v2").paginate(Bucket=self.bucket):
            total += sum(obj["Size"] for obj in page.get("Contents", []))
        return total

    def set_expiry(self, days: int, prefixes: tuple[str, ...] = ("runs/", "jobs/")) -> None:
        """Delete run and job files automatically after `days` (keeps the bucket inside R2's free 10 GB)."""
        self._s3.put_bucket_lifecycle_configuration(
            Bucket=self.bucket,
            LifecycleConfiguration={
                "Rules": [
                    {
                        "ID": f"expire-{p.strip('/')}",
                        "Status": "Enabled",
                        "Filter": {"Prefix": p},
                        "Expiration": {"Days": days},
                    }
                    for p in prefixes
                ]
            },
        )

    def set_cors(self, origins: list[str]) -> None:
        """Let the Foxglove web app (and the report page) fetch run files directly from the bucket."""
        self._s3.put_bucket_cors(
            Bucket=self.bucket,
            CORSConfiguration={
                "CORSRules": [
                    {
                        "AllowedOrigins": origins,
                        "AllowedMethods": ["GET", "HEAD"],
                        "AllowedHeaders": ["*"],
                        "ExposeHeaders": ["Content-Length", "Content-Range", "Accept-Ranges", "ETag"],
                        "MaxAgeSeconds": 3600,
                    }
                ]
            },
        )


class LocalStorage:
    """Directory-backed stand-in for R2 (tests and offline dry runs)."""

    def __init__(self, root: Path, base_url: str = "file://"):
        self.root = Path(root)
        self.base_url = base_url

    def upload(self, local: Path, key: str) -> None:
        dest = self.root / key
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(Path(local).read_bytes())

    def url(self, key: str, expires_s: int = 0) -> str:
        return (
            f"{self.base_url}{(self.root / key).resolve()}" if self.base_url == "file://" else f"{self.base_url}/{key}"
        )

    def read_text(self, key: str) -> str | None:
        p = self.root / key
        return p.read_text() if p.is_file() else None

    def total_bytes(self) -> int:
        return sum(p.stat().st_size for p in self.root.rglob("*") if p.is_file()) if self.root.exists() else 0


def run_prefix(run_id: str) -> str:
    return f"runs/{run_id}"


def main(argv: list[str] | None = None) -> int:
    """python -m api.storage cors | check"""
    import argparse
    import sys
    import tempfile

    parser = argparse.ArgumentParser(description="Cloudflare R2 setup for simulation runs")
    sub = parser.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("cors", help="allow the Foxglove web app (and extra origins) to read run files")
    p.add_argument("--origin", action="append", default=[], help="extra allowed origin, e.g. your API's URL")
    sub.add_parser("check", help="upload a test object and print a presigned link")
    sub.add_parser("usage", help="how much the bucket holds")
    p = sub.add_parser("expire", help="delete run/job files automatically after N days")
    p.add_argument("--days", type=int, default=30)
    args = parser.parse_args(argv)
    config = R2Config.from_env()
    if not config:
        print("error: set R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY and R2_BUCKET", file=sys.stderr)
        return 1
    r2 = R2Storage(config)
    if args.cmd == "cors":
        # embed.foxglove.dev: the viewer embedded in the report page; app.foxglove.dev: the full web app.
        origins = ["https://embed.foxglove.dev", "https://app.foxglove.dev", *args.origin]
        r2.set_cors(origins)
        print(f"CORS on bucket {config.bucket}: GET/HEAD from {', '.join(origins)}")
        return 0
    if args.cmd == "usage":
        print(f"bucket {config.bucket}: {r2.total_bytes() / 1e9:.2f} GB")
        return 0
    if args.cmd == "expire":
        r2.set_expiry(args.days)
        print(f"bucket {config.bucket}: files under runs/ and jobs/ are deleted after {args.days} days")
        return 0
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
        f.write('{"ok": true}\n')
    r2.upload(Path(f.name), "runs/_check/check.json")
    print(f"uploaded runs/_check/check.json; presigned link (7 days):\n{r2.url('runs/_check/check.json')}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())