diff --git a/README.md b/README.md index d2bb6b3..6d2050f 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,16 @@ qc-cli infra destroy --yes Destroy stack without confirmation qc-cli infra destroy --delete-bucket-data Destroy stack and delete S3 data ``` +### `upload` + +``` +qc-cli upload Upload a single file to S3 +qc-cli upload Upload all files in a directory tree to S3 +qc-cli upload --s3-key Upload a file to a custom S3 key +``` + +Uploads use `s3.bucket` and `s3.data_prefix` from `config.yaml`. File uploads default to `s3:////`. Directory uploads are recursive, preserve paths relative to the uploaded directory, and place files under `s3:////`. + ## AWS permissions required The IAM user or role running the CLI needs: diff --git a/src/aws/s3.py b/src/aws/s3.py new file mode 100644 index 0000000..9a2141a --- /dev/null +++ b/src/aws/s3.py @@ -0,0 +1,51 @@ +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +import boto3 +from mypy_boto3_s3 import S3Client + + +def _client(region: str, profile: str) -> S3Client: + return boto3.Session(profile_name=profile, region_name=region).client("s3") + + +def upload_file( + region: str, + profile: str, + bucket: str, + local_path: str, + s3_key: str, +) -> str: + _client(region, profile).upload_file(local_path, bucket, s3_key) + return f"s3://{bucket}/{s3_key}" + + +def upload_dir( + region: str, + profile: str, + bucket: str, + local_dir: str, + s3_prefix: str, + on_progress: Callable[[], None] | None = None, +) -> int: + root = Path(local_dir) + files = [file for file in root.rglob("*") if file.is_file()] + if not files: + return 0 + + client = _client(region, profile) + prefix = s3_prefix.rstrip("/") + + def upload_one(file_path: Path) -> None: + key = f"{prefix}/{file_path.relative_to(root)}" + client.upload_file(str(file_path), bucket, key) + if on_progress: + on_progress() + + with ThreadPoolExecutor(max_workers=10) as pool: + futures = [pool.submit(upload_one, file) for file in files] + for future in as_completed(futures): + future.result() + + return len(files) diff --git a/src/main.py b/src/main.py index dd8cd9a..3414744 100644 --- a/src/main.py +++ b/src/main.py @@ -3,8 +3,11 @@ from pathlib import Path import typer import yaml from rich.console import Console +from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn +from src.aws import s3 as s3_ops from src.commands import infra +from src.commands.utils import CONFIG_OPT, load_cfg from src.config import Config app = typer.Typer( @@ -34,3 +37,64 @@ def init( console.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]") console.print("Edit it (especially [cyan]s3.bucket[/cyan]) before running other commands.") + + +@app.command() +def upload( + path: Path = typer.Argument(..., help="Local file or directory to upload"), + s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"), + config: str = CONFIG_OPT, +) -> None: + """Upload a local file or directory to S3.""" + cfg = load_cfg(config) + + if path.is_file(): + key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}" + try: + with console.status(f"Uploading {path.name}..."): + uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key) + except Exception as e: + console.print(f"[red]Upload failed: {e}[/red]") + raise typer.Exit(1) + + console.print(f"[green]✓[/green] {path.name} -> {uri}") + return + + if path.is_dir(): + if s3_key is not None: + console.print("[red]--s3-key can only be used when uploading a single file.[/red]") + raise typer.Exit(1) + + files = [file for file in path.rglob("*") if file.is_file()] + if not files: + console.print("[yellow]No files found in directory.[/yellow]") + raise typer.Exit(0) + + prefix = cfg.s3.data_prefix + console.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/") + try: + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + console=console, + ) as progress: + task = progress.add_task("Uploading...", total=len(files)) + count = s3_ops.upload_dir( + cfg.aws.region, + cfg.aws.profile, + cfg.s3.bucket, + str(path), + prefix, + on_progress=lambda: progress.advance(task), + ) + except Exception as e: + console.print(f"[red]Upload failed: {e}[/red]") + raise typer.Exit(1) + + console.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/") + return + + console.print(f"[red]Path not found: {path}[/red]") + raise typer.Exit(1)