Streamlit_NBA/scripts/db_storage.py

#!/usr/bin/env python3
"""
S3/Cloudflare R2 Object Storage utility for DuckDB database management.

Supports:
- Cloudflare R2 (recommended: zero egress fees)
- Hetzner S3 Object Storage
- SeaweedFS / MinIO / AWS S3

Commands:
  python scripts/db_storage.py status
  python scripts/db_storage.py upload [--file PATH] [--key KEY]
  python scripts/db_storage.py download [--force] [--key KEY]
"""

import argparse
import hashlib
import os
import shutil
import sys
import time
from pathlib import Path

try:
    from dotenv import load_dotenv
    # Load .env file from repo root
    repo_root = Path(__file__).resolve().parent.parent
    load_dotenv(repo_root / ".env")
except ImportError:
    repo_root = Path(__file__).resolve().parent.parent

try:
    import boto3
    from boto3.s3.transfer import TransferConfig
    from botocore.client import Config
    from botocore.exceptions import ClientError
except ImportError:
    print("ERROR: boto3 is not installed. Please run `uv sync` to install dependencies.", file=sys.stderr)
    sys.exit(1)

# Default Paths
DEFAULT_LOCAL_DB = repo_root / "data" / "DB" / "dbt_nba.duckdb"
DEFAULT_REPORTS_DB = repo_root / "dbt_nba" / "reports" / "sources" / "nba" / "dbt_nba.duckdb"

# Tuned TransferConfig for resilient large multipart transfers over HTTPS
TRANSFER_CONFIG = TransferConfig(
    multipart_threshold=32 * 1024 * 1024,   # 32 MB
    max_concurrency=4,                       # 4 concurrent threads
    multipart_chunksize=32 * 1024 * 1024,   # 32 MB per chunk
    use_threads=True,
    num_download_attempts=5,
)


def get_s3_config():
    """Retrieve S3/R2 configuration from environment variables."""
    endpoint_url = os.getenv("S3_ENDPOINT_URL") or os.getenv("R2_ENDPOINT_URL")
    bucket_name = os.getenv("S3_BUCKET_NAME") or os.getenv("R2_BUCKET_NAME")
    access_key = os.getenv("S3_ACCESS_KEY_ID") or os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("R2_ACCESS_KEY_ID")
    secret_key = os.getenv("S3_SECRET_ACCESS_KEY") or os.getenv("AWS_SECRET_ACCESS_KEY") or os.getenv("R2_SECRET_ACCESS_KEY")
    region_name = os.getenv("S3_REGION", "auto")
    db_key = os.getenv("S3_DB_KEY", "dbt_nba.duckdb")

    missing = []
    if not endpoint_url:
        missing.append("S3_ENDPOINT_URL")
    if not bucket_name:
        missing.append("S3_BUCKET_NAME")
    if not access_key:
        missing.append("S3_ACCESS_KEY_ID")
    if not secret_key:
        missing.append("S3_SECRET_ACCESS_KEY")

    if missing:
        print(f"ERROR: Missing required environment variables in .env: {', '.join(missing)}", file=sys.stderr)
        print("Please check .env.example for guidance.", file=sys.stderr)
        sys.exit(1)

    return {
        "endpoint_url": endpoint_url,
        "bucket_name": bucket_name,
        "access_key": access_key,
        "secret_key": secret_key,
        "region_name": region_name,
        "db_key": db_key,
    }


def get_s3_client(config):
    """Create a boto3 S3 client configured for custom S3/R2 endpoints."""
    return boto3.client(
        "s3",
        endpoint_url=config["endpoint_url"],
        aws_access_key_id=config["access_key"],
        aws_secret_access_key=config["secret_key"],
        region_name=config["region_name"],
        config=Config(
            signature_version="s3v4",
            s3={"addressing_style": "path"},
            retries={"max_attempts": 5, "mode": "adaptive"},
        ),
    )


class ProgressPercentage:
    """Console progress tracker for S3 file transfer."""

    def __init__(self, filename, total_size, action="Transferring"):
        self._filename = filename
        self._total_size = float(total_size) if total_size > 0 else 1.0
        self._seen_so_far = 0
        self._start_time = time.time()
        self._action = action

    def __call__(self, bytes_amount):
        self._seen_so_far += bytes_amount
        percentage = (self._seen_so_far / self._total_size) * 100
        elapsed = max(time.time() - self._start_time, 0.001)
        speed_mb = (self._seen_so_far / (1024 * 1024)) / elapsed
        mb_seen = self._seen_so_far / (1024 * 1024)
        mb_total = self._total_size / (1024 * 1024)

        sys.stdout.write(
            f"\r{self._action} {Path(self._filename).name}: {mb_seen:.1f}/{mb_total:.1f} MB "
            f"({percentage:5.1f}%) @ {speed_mb:5.1f} MB/s"
        )
        sys.stdout.flush()
        if self._seen_so_far >= self._total_size:
            sys.stdout.write("\n")
            sys.stdout.flush()


def compute_file_md5(filepath: Path) -> str:
    """Compute hex MD5 hash of a local file."""
    hash_md5 = hashlib.md5()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def upload_db(file_path: Path = None, key: str = None):
    """Upload DuckDB file to S3/R2 with auto-retry."""
    config = get_s3_config()
    target_file = Path(file_path) if file_path else DEFAULT_LOCAL_DB
    object_key = key or config["db_key"]

    if not target_file.exists():
        print(f"ERROR: Local database file not found at {target_file}", file=sys.stderr)
        sys.exit(1)

    file_size = target_file.stat().st_size
    print(f"==> Uploading {target_file} ({file_size / (1024*1024):.2f} MB) to s3://{config['bucket_name']}/{object_key}...")

    s3 = get_s3_client(config)
    max_retries = 3

    for attempt in range(1, max_retries + 1):
        try:
            progress = ProgressPercentage(str(target_file), file_size, action="Uploading")
            s3.upload_file(
                Filename=str(target_file),
                Bucket=config["bucket_name"],
                Key=object_key,
                Config=TRANSFER_CONFIG,
                Callback=progress,
            )
            print(f"✓ Successfully uploaded {object_key} to {config['bucket_name']}")
            return
        except Exception as e:
            if attempt < max_retries:
                print(f"\n[Attempt {attempt}/{max_retries}] Upload encountered error ({e}). Retrying in 3s...", file=sys.stderr)
                time.sleep(3)
            else:
                print(f"\nERROR: Failed to upload to S3 after {max_retries} attempts: {e}", file=sys.stderr)
                sys.exit(1)


def download_db(force: bool = False, key: str = None):
    """Download DuckDB file from S3/R2 to local destinations with auto-retry."""
    config = get_s3_config()
    object_key = key or config["db_key"]
    s3 = get_s3_client(config)

    print(f"==> Checking remote object s3://{config['bucket_name']}/{object_key}...")
    try:
        head = s3.head_object(Bucket=config["bucket_name"], Key=object_key)
    except ClientError as e:
        print(f"ERROR: Could not find or access remote object {object_key} in {config['bucket_name']}: {e}", file=sys.stderr)
        sys.exit(1)

    remote_size = head.get("ContentLength", 0)
    last_modified = head.get("LastModified")

    print(f"    Remote size: {remote_size / (1024*1024):.2f} MB | Last modified: {last_modified}")

    # Check if local file is already identical
    if not force and DEFAULT_LOCAL_DB.exists() and DEFAULT_LOCAL_DB.stat().st_size == remote_size:
        # Check if reports db also exists
        if DEFAULT_REPORTS_DB.exists() and DEFAULT_REPORTS_DB.stat().st_size == remote_size:
            print("✓ Local database is already up-to-date. Skipping download (use --force to re-download).")
            return

    # Ensure parent directories exist
    DEFAULT_LOCAL_DB.parent.mkdir(parents=True, exist_ok=True)
    DEFAULT_REPORTS_DB.parent.mkdir(parents=True, exist_ok=True)

    temp_file = DEFAULT_LOCAL_DB.with_suffix(".duckdb.tmp")
    max_retries = 3

    for attempt in range(1, max_retries + 1):
        try:
            progress = ProgressPercentage(str(DEFAULT_LOCAL_DB), remote_size, action="Downloading")
            s3.download_file(
                Bucket=config["bucket_name"],
                Key=object_key,
                Filename=str(temp_file),
                Config=TRANSFER_CONFIG,
                Callback=progress,
            )
            # Atomically move temp file to main db path
            shutil.move(str(temp_file), str(DEFAULT_LOCAL_DB))
            # Copy to reports location as well
            shutil.copy2(str(DEFAULT_LOCAL_DB), str(DEFAULT_REPORTS_DB))
            print(f"✓ Downloaded database to:\n  - {DEFAULT_LOCAL_DB}\n  - {DEFAULT_REPORTS_DB}")
            return
        except Exception as e:
            if temp_file.exists():
                temp_file.unlink()
            if attempt < max_retries:
                print(f"\n[Attempt {attempt}/{max_retries}] Download encountered error ({e}). Retrying in 3s...", file=sys.stderr)
                time.sleep(3)
            else:
                print(f"\nERROR: Failed to download from S3 after {max_retries} attempts: {e}", file=sys.stderr)
                sys.exit(1)


def status_db():
    """Show status of remote and local databases."""
    config = get_s3_config()
    object_key = config["db_key"]
    s3 = get_s3_client(config)

    print("=" * 60)
    print("Database Storage Status (Cloudflare R2 / S3)")
    print("=" * 60)
    print(f"Endpoint: {config['endpoint_url']}")
    print(f"Bucket:   {config['bucket_name']}")
    print(f"DB Key:   {object_key}")
    print("-" * 60)

    try:
        head = s3.head_object(Bucket=config["bucket_name"], Key=object_key)
        remote_size = head.get("ContentLength", 0)
        last_modified = head.get("LastModified")
        print(f"Remote DB: Exists ({remote_size / (1024*1024):.2f} MB)")
        print(f"           Last Modified: {last_modified}")
    except ClientError as e:
        print(f"Remote DB: NOT FOUND ({e})")

    print("-" * 60)
    if DEFAULT_LOCAL_DB.exists():
        print(f"Local DB (Data):    Exists ({DEFAULT_LOCAL_DB.stat().st_size / (1024*1024):.2f} MB) -> {DEFAULT_LOCAL_DB}")
    else:
        print(f"Local DB (Data):    Missing -> {DEFAULT_LOCAL_DB}")

    if DEFAULT_REPORTS_DB.exists():
        print(f"Local DB (Reports): Exists ({DEFAULT_REPORTS_DB.stat().st_size / (1024*1024):.2f} MB) -> {DEFAULT_REPORTS_DB}")
    else:
        print(f"Local DB (Reports): Missing -> {DEFAULT_REPORTS_DB}")
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(description="Manage NBA DuckDB sync with Cloudflare R2 / S3.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # status
    subparsers.add_parser("status", help="Check remote and local database status")

    # upload
    upload_parser = subparsers.add_parser("upload", help="Upload local DuckDB to S3/R2")
    upload_parser.add_argument("--file", help="Path to local .duckdb file (default: data/DB/dbt_nba.duckdb)")
    upload_parser.add_argument("--key", help="Remote object key (default: from env or dbt_nba.duckdb)")

    # download
    download_parser = subparsers.add_parser("download", help="Download DuckDB from S3/R2")
    download_parser.add_argument("--force", action="store_true", help="Force download even if local file matches")
    download_parser.add_argument("--key", help="Remote object key (default: from env or dbt_nba.duckdb)")

    args = parser.parse_args()

    if args.command == "status":
        status_db()
    elif args.command == "upload":
        upload_db(file_path=args.file, key=args.key)
    elif args.command == "download":
        download_db(force=args.force, key=args.key)


if __name__ == "__main__":
    main()