Downloading files

"""Download a partition from Cloud Humans Data Export."""

from __future__ import annotations

from pathlib import Path

import boto3
from boto3.s3.transfer import TransferConfig
from botocore.exceptions import ClientError

BUCKET = "<YOUR_BUCKET>"
# None on AWS compute, so the SDK uses the role the environment already provides.
PROFILE = "cloudhumans-export"

# 20 threads instead of the default 10. Beyond that the gain flattens out and the risk of 503s increases.
TRANSFER = TransferConfig(max_concurrency=20, multipart_threshold=16 * 1024 * 1024)


class PartitionNotReady(Exception):
    """The partition exists but doesn't have the _SUCCESS marker yet."""


class ObjectArchived(Exception):
    """The object is in Glacier Flexible Retrieval and requires a restore before it can be read."""


def _client():
    session = boto3.Session(profile_name=PROFILE) if PROFILE else boto3.Session()
    return session.client("s3")


def partition_prefix(table: str, partition: str) -> str:
    """`partition` is the folder name, `snapshot_date=2026-08-06` or `etl_date=2026-08-06`."""
    return f"v1/datamart_claudia/{table}/{partition}/"


def is_ready(s3, prefix: str) -> bool:
    try:
        s3.head_object(Bucket=BUCKET, Key=f"{prefix}_SUCCESS")
        return True
    except ClientError as err:
        if err.response["Error"]["Code"] in ("404", "NoSuchKey"):
            return False
        raise


def download_partition(table: str, partition: str, dest: Path) -> list[Path]:
    s3 = _client()
    prefix = partition_prefix(table, partition)

    if not is_ready(s3, prefix):
        raise PartitionNotReady(prefix)

    dest.mkdir(parents=True, exist_ok=True)
    written: list[Path] = []

    for page in s3.get_paginator("list_objects_v2").paginate(Bucket=BUCKET, Prefix=prefix):
        for obj in page.get("Contents", []):
            key = obj["Key"]
            if key.endswith("/") or key.endswith("_SUCCESS"):
                continue

            target = dest / Path(key).name
            try:
                s3.download_file(BUCKET, key, str(target), Config=TRANSFER)
            except ClientError as err:
                code = err.response["Error"]["Code"]
                if code == "InvalidObjectState":
                    raise ObjectArchived(key) from err
                raise
            written.append(target)

    return written


if __name__ == "__main__":
    files = download_partition("<TABLE>", "snapshot_date=<DATE>", Path("./data"))
    print(f"{len(files)} files downloaded")