Executable
+302
@@ -0,0 +1,302 @@
|
||||
#!/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()
|
||||
@@ -0,0 +1,28 @@
|
||||
-- Extract source tables from Postgres into DuckDB
|
||||
-- Requires: POSTGRES_URL environment variable
|
||||
-- Usage: Run via scripts/pipeline.sh (handles variable substitution)
|
||||
|
||||
INSTALL postgres;
|
||||
LOAD postgres;
|
||||
|
||||
ATTACH '${POSTGRES_URL}' AS pg (TYPE POSTGRES, READ_ONLY);
|
||||
|
||||
DROP TABLE IF EXISTS main.games;
|
||||
DROP TABLE IF EXISTS main.line_scores;
|
||||
DROP TABLE IF EXISTS main.player_game_basic_stats;
|
||||
DROP TABLE IF EXISTS main.player_game_adv_stats;
|
||||
DROP TABLE IF EXISTS main.player_shot_charts;
|
||||
DROP TABLE IF EXISTS main.team_game_basic_stats;
|
||||
DROP TABLE IF EXISTS main.team_game_adv_stats;
|
||||
|
||||
CREATE TABLE main.games AS SELECT * FROM pg.public.games;
|
||||
CREATE TABLE main.line_scores AS SELECT * FROM pg.public.line_scores;
|
||||
CREATE TABLE main.player_game_basic_stats AS SELECT * FROM pg.public.player_game_basic_stats;
|
||||
CREATE TABLE main.player_game_adv_stats AS SELECT * FROM pg.public.player_game_adv_stats;
|
||||
CREATE TABLE main.player_shot_charts AS SELECT * FROM pg.public.player_shot_charts;
|
||||
CREATE TABLE main.team_game_basic_stats AS SELECT * FROM pg.public.team_game_basic_stats;
|
||||
CREATE TABLE main.team_game_adv_stats AS SELECT * FROM pg.public.team_game_adv_stats;
|
||||
|
||||
DETACH pg;
|
||||
|
||||
SELECT 'Extraction complete' AS status;
|
||||
Executable
+47
@@ -0,0 +1,47 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
DB_PATH="$REPO_ROOT/data/DB/dbt_nba.duckdb"
|
||||
REPORTS_DB="$REPO_ROOT/dbt_nba/reports/sources/nba/dbt_nba.duckdb"
|
||||
|
||||
# Load .env if it exists
|
||||
if [ -f "$REPO_ROOT/.env" ]; then
|
||||
set -a
|
||||
source "$REPO_ROOT/.env"
|
||||
set +a
|
||||
fi
|
||||
|
||||
# Check prerequisites
|
||||
if [ -z "${POSTGRES_URL:-}" ]; then
|
||||
echo "ERROR: POSTGRES_URL environment variable is not set" >&2
|
||||
echo "Please export POSTGRES_URL or configure it in $REPO_ROOT/.env" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "==> [1/4] Extracting from Postgres into DuckDB..."
|
||||
mkdir -p "$(dirname "$DB_PATH")"
|
||||
mkdir -p "$(dirname "$REPORTS_DB")"
|
||||
rm -f "$DB_PATH"
|
||||
envsubst < "$REPO_ROOT/scripts/extract.sql" | duckdb "$DB_PATH"
|
||||
|
||||
echo "==> [2/4] Running dbt build --full-refresh..."
|
||||
export DBT_PROJECT_DIR="$REPO_ROOT/dbt_nba"
|
||||
export DBT_PROFILES_DIR="$REPO_ROOT/dbt_nba"
|
||||
export DBT_DUCKDB_PATH="$DB_PATH"
|
||||
|
||||
uv run --project "$REPO_ROOT" dbt deps --project-dir "$DBT_PROJECT_DIR" --profiles-dir "$DBT_PROFILES_DIR"
|
||||
uv run --project "$REPO_ROOT" dbt build --full-refresh --project-dir "$DBT_PROJECT_DIR" --profiles-dir "$DBT_PROFILES_DIR"
|
||||
|
||||
echo "==> [3/4] Copying database to reports directory..."
|
||||
cp "$DB_PATH" "$REPORTS_DB"
|
||||
|
||||
echo "==> [4/4] Syncing database artifact to Cloudflare R2 / S3..."
|
||||
if [ -n "${S3_ENDPOINT_URL:-}" ] || [ -n "${R2_ENDPOINT_URL:-}" ]; then
|
||||
uv run --project "$REPO_ROOT" python "$REPO_ROOT/scripts/db_storage.py" upload
|
||||
echo "==> Cloudflare R2 upload complete!"
|
||||
else
|
||||
echo "==> S3_ENDPOINT_URL not set; skipping remote upload (local DB ready)."
|
||||
fi
|
||||
|
||||
echo "==> Pipeline complete successfully!"
|
||||
Executable
+41
@@ -0,0 +1,41 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
REPORTS_DB="$REPO_ROOT/dbt_nba/reports/sources/nba/dbt_nba.duckdb"
|
||||
DATA_DB="$REPO_ROOT/data/DB/dbt_nba.duckdb"
|
||||
|
||||
# Load .env if present
|
||||
if [ -f "$REPO_ROOT/.env" ]; then
|
||||
set -a
|
||||
source "$REPO_ROOT/.env"
|
||||
set +a
|
||||
fi
|
||||
|
||||
FORCE_SYNC=0
|
||||
for arg in "$@"; do
|
||||
if [ "$arg" == "--sync" ] || [ "$arg" == "-s" ]; then
|
||||
FORCE_SYNC=1
|
||||
fi
|
||||
done
|
||||
|
||||
# Check if database is missing or sync requested
|
||||
if [ ! -f "$REPORTS_DB" ] || [ "$FORCE_SYNC" -eq 1 ]; then
|
||||
if [ -f "$DATA_DB" ] && [ "$FORCE_SYNC" -eq 0 ]; then
|
||||
echo "==> Copying local database to reports path..."
|
||||
mkdir -p "$(dirname "$REPORTS_DB")"
|
||||
cp "$DATA_DB" "$REPORTS_DB"
|
||||
elif [ -n "${S3_ENDPOINT_URL:-}" ] || [ -n "${R2_ENDPOINT_URL:-}" ]; then
|
||||
echo "==> Local database missing or --sync requested; downloading latest from Cloudflare R2 / S3..."
|
||||
uv run --project "$REPO_ROOT" python "$REPO_ROOT/scripts/db_storage.py" download
|
||||
else
|
||||
echo "ERROR: Local database file not found at $REPORTS_DB." >&2
|
||||
echo "Please either:" >&2
|
||||
echo " 1. Run ./scripts/pipeline.sh to extract and build from Postgres, or" >&2
|
||||
echo " 2. Configure R2 credentials in .env and run: uv run python scripts/db_storage.py download" >&2
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
echo "==> Starting Streamlit NBA Analytics Dashboard..."
|
||||
uv run --project "$REPO_ROOT" streamlit run "$REPO_ROOT/streamlit_app/app.py"
|
||||
Reference in New Issue
Block a user