social-data-attribution-w2/scripts/r2_sunset/delete_prefix.py

"""Delete an R2 prefix (or sub-prefix) under the soc127-dedup bucket.

Safety:
- Defaults to --dry-run: lists what WOULD be deleted, no mutation.
- --execute additionally requires the env var R2_DELETE_OK=1 in the same
  invocation. Without both, nothing is deleted.
- Re-lists the prefix from R2 live (not the stale crawl) so it deletes exactly
  current reality, and re-lists afterward to confirm zero objects remain.

Deletion uses batched delete_objects (up to 1000 keys/request); Class A cost is
negligible (~$0.00027 for 60k objects).

Run via:
    bash scripts/bootstrap/with_r2_credentials.sh \\
        uv run python scripts/r2_sunset/delete_prefix.py --prefix soc91-stats/ --dry-run
    R2_DELETE_OK=1 bash scripts/bootstrap/with_r2_credentials.sh \\
        uv run python scripts/r2_sunset/delete_prefix.py --prefix soc91-stats/ --execute
"""

from __future__ import annotations

import argparse
import logging
import os
import sys

import boto3

R2_BUCKET = "soc127-dedup"
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"

logger = logging.getLogger("delete_prefix")


def list_prefix(s3, prefix: str) -> list[tuple[str, int]]:
    out: list[tuple[str, int]] = []
    paginator = s3.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=R2_BUCKET, Prefix=prefix):
        for obj in page.get("Contents", []):
            out.append((obj["Key"], obj["Size"]))
    return out


def fmt_gb(b: int) -> str:
    return f"{b / 1024**3:.2f} GB"


def delete_in_batches(s3, keys: list[str]) -> int:
    deleted = 0
    for i in range(0, len(keys), 1000):
        batch = [{"Key": k} for k in keys[i : i + 1000]]
        resp = s3.delete_objects(
            Bucket=R2_BUCKET, Delete={"Objects": batch, "Quiet": True}
        )
        errors = resp.get("Errors", [])
        if errors:
            for e in errors[:5]:
                logger.error("delete error: %s", e)
            raise SystemExit(f"delete_objects reported {len(errors)} errors; aborting")
        deleted += len(batch)
        logger.info("  deleted %d/%d", deleted, len(keys))
    return deleted


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--prefix", required=True, help="R2 key prefix to delete (e.g. soc91-stats/)"
    )
    mode = ap.add_mutually_exclusive_group()
    mode.add_argument("--dry-run", action="store_true", default=True)
    mode.add_argument("--execute", action="store_true")
    ap.add_argument("--log-level", default="INFO")
    args = ap.parse_args()

    logging.basicConfig(
        level=args.log_level, format="%(asctime)s %(levelname)s %(message)s"
    )

    if not os.environ.get("R2_ACCESS_KEY_ID") or not os.environ.get(
        "R2_SECRET_ACCESS_KEY"
    ):
        logger.error(
            "R2 creds missing. Run via: bash scripts/bootstrap/with_r2_credentials.sh "
            "uv run python %s ...",
            sys.argv[0],
        )
        return 2

    s3 = boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT_URL,
        aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
    )

    logger.info("listing r2://%s/%s ...", R2_BUCKET, args.prefix)
    objs = list_prefix(s3, args.prefix)
    total_bytes = sum(s for _, s in objs)
    logger.info("found %d objects, %s", len(objs), fmt_gb(total_bytes))
    for k, _ in objs[:5]:
        logger.info("  sample: %s", k)

    if not objs:
        logger.info("nothing to delete under %s", args.prefix)
        return 0

    execute = args.execute and os.environ.get("R2_DELETE_OK") == "1"
    if not args.execute:
        logger.info(
            "DRY-RUN: would delete %d objects (%s). No mutation.",
            len(objs),
            fmt_gb(total_bytes),
        )
        return 0
    if not execute:
        logger.error(
            "--execute given but R2_DELETE_OK!=1. Refusing to delete. "
            "Set R2_DELETE_OK=1 in the same command to confirm."
        )
        return 3

    logger.warning(
        "EXECUTING deletion of %d objects (%s) under %s",
        len(objs),
        fmt_gb(total_bytes),
        args.prefix,
    )
    deleted = delete_in_batches(s3, [k for k, _ in objs])

    remaining = list_prefix(s3, args.prefix)
    logger.info(
        "deleted %d objects; %d remain under %s", deleted, len(remaining), args.prefix
    )
    if remaining:
        logger.warning(
            "non-empty after delete (%d left); inspect manually", len(remaining)
        )
        return 4
    logger.info("OK: prefix %s now empty. Freed %s.", args.prefix, fmt_gb(total_bytes))
    return 0


if __name__ == "__main__":
    sys.exit(main())