scripts/r2_storage.py

"""Cloudflare R2 (S3-compatible) helpers."""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import boto3
from botocore.client import BaseClient
from botocore.config import Config

DEFAULT_BUCKET = "antibionlp-label-studio"
PREFIX_LLM = "llm"
PREFIX_TASKS = "tasks"
PREFIX_LS = "label-studio"


def load_r2_env(env: dict[str, str] | None = None) -> dict[str, str]:
    src = env if env is not None else os.environ
    access = src.get("R2_ACCESS_KEY_ID") or src.get("AWS_ACCESS_KEY_ID")
    secret = src.get("R2_SECRET_ACCESS_KEY") or src.get("AWS_SECRET_ACCESS_KEY")
    endpoint = src.get("R2_ENDPOINT_URL") or src.get("AWS_ENDPOINT_URL")
    bucket = src.get("R2_BUCKET_NAME", DEFAULT_BUCKET)
    if not access or not secret or not endpoint:
        raise RuntimeError(
            "R2 credentials missing. Set R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, "
            "R2_ENDPOINT_URL (and optionally R2_BUCKET_NAME)."
        )
    return {
        "access_key": access,
        "secret_key": secret,
        "endpoint_url": endpoint.rstrip("/"),
        "bucket": bucket,
    }


def r2_client(cfg: dict[str, str] | None = None) -> BaseClient:
    c = load_r2_env(cfg)
    return boto3.client(
        "s3",
        endpoint_url=c["endpoint_url"],
        aws_access_key_id=c["access_key"],
        aws_secret_access_key=c["secret_key"],
        region_name="auto",
        config=Config(signature_version="s3v4", s3={"addressing_style": "path"}),
    )


def bucket_name(cfg: dict[str, str] | None = None) -> str:
    return load_r2_env(cfg)["bucket"]


def download_prefix(
    prefix: str,
    local_dir: Path,
    *,
    client: BaseClient | None = None,
    cfg: dict[str, str] | None = None,
) -> int:
    c = client or r2_client(cfg)
    b = bucket_name(cfg)
    local_dir.mkdir(parents=True, exist_ok=True)
    paginator = c.get_paginator("list_objects_v2")
    n = 0
    base = prefix.rstrip("/") + "/"
    for page in paginator.paginate(Bucket=b, Prefix=base):
        for obj in page.get("Contents") or []:
            key = obj["Key"]
            if key == base or not key.startswith(base):
                continue
            rel = key[len(base) :]
            if not rel:
                continue
            dest = local_dir / rel
            dest.parent.mkdir(parents=True, exist_ok=True)
            c.download_file(b, key, str(dest))
            n += 1
    return n


def sync_directory(
    local_dir: Path,
    prefix: str,
    *,
    client: BaseClient | None = None,
    cfg: dict[str, str] | None = None,
) -> int:
    if not local_dir.is_dir():
        return 0
    c = client or r2_client(cfg)
    b = bucket_name(cfg)
    root = prefix.rstrip("/")
    n = 0
    for path in local_dir.rglob("*"):
        if not path.is_file():
            continue
        key = f"{root}/{path.relative_to(local_dir).as_posix()}"
        c.upload_file(str(path), b, key)
        n += 1
    return n