52 lines
1.3 KiB
Python
52 lines
1.3 KiB
Python
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)
|