Source code for

"""Service to upload test objects to a managed bucket."""

from pathlib import Path

import boto3
import botocore
import click

from latch.utils import account_id_from_token, retrieve_or_login
from import _retrieve_creds

BUCKET = "latch-public"

[docs]def upload(src_path: str, dont_confirm_overwrite: bool = True) -> str: """Uploads a local file/folder to a managed bucket. Args: src_path: The path of the file/folder to upload. Returns: s3 URL of uploaded object. Example: >>> upload("./foo.txt") # Puts the file `foo.txt` in a managed bucket and returns a publicly # accesible S3 URL. """ session_token, access_key, secret_key, account_id = _retrieve_creds() s3_resource = boto3.client( "s3", aws_access_key_id=access_key, aws_secret_access_key=secret_key, aws_session_token=session_token, ) if account_id is None or account_id == "": account_id = account_id_from_token(retrieve_or_login()) allowed_key = str((Path("test-data") / account_id).joinpath(src_path)) upload_helper( Path(src_path).resolve(), allowed_key, s3_resource, dont_confirm_overwrite, ) return f"s3://{BUCKET}/{allowed_key}"
[docs]def upload_helper( src_path: Path, key: str, s3_resource, dont_confirm_overwrite: bool, ): if not src_path.exists(): raise ValueError(f"Path {src_path} doesn't exist.") if src_path.is_dir(): for sub_path in src_path.iterdir(): upload_helper( sub_path, f"{key}/{}", s3_resource, dont_confirm_overwrite ) else: upload_file(src_path, key, s3_resource, dont_confirm_overwrite)
[docs]def upload_file( src_path: Path, key: str, s3_resource, dont_confirm_overwrite: bool, ): try: s3_resource.head_object(Bucket=BUCKET, Key=key) except botocore.exceptions.ClientError: pass else: while True and not dont_confirm_overwrite: confirm = input(f"{key} already exists, override it? (y/n) > ") if confirm in ("n", "no"): print("Aborting upload.") return elif confirm in ("y", "yes"): break else: print("Invalid response.") s3_resource.upload_file(str(src_path), BUCKET, key) click.secho(f"Successfully uploaded {key}", fg="green")