kleinkram 0.38.1.dev20241119134715__py3-none-any.whl → 0.38.1.dev20241125112529__py3-none-any.whl

kleinkram/api/file_transfer.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
-import os
+import logging
 import sys
+from concurrent.futures import as_completed
+from concurrent.futures import Future
 from concurrent.futures import ThreadPoolExecutor
+from enum import Enum
 from pathlib import Path
 from time import monotonic
 from typing import Dict
-from typing import List
 from typing import NamedTuple
 from typing import Optional
 from typing import Tuple
@@ -18,14 +20,19 @@ import httpx
 from kleinkram.api.client import AuthenticatedClient
 from kleinkram.config import Config
 from kleinkram.config import LOCAL_S3
-from kleinkram.errors import AccessDeniedException
-from kleinkram.errors import CorruptedFile
-from kleinkram.errors import UploadFailed
+from kleinkram.errors import AccessDenied
+from kleinkram.models import File
+from kleinkram.models import FileState
 from kleinkram.utils import b64_md5
-from kleinkram.utils import raw_rich
-from rich.text import Text
+from kleinkram.utils import format_error
+from kleinkram.utils import format_traceback
+from kleinkram.utils import styled_string
+from rich.console import Console
 from tqdm import tqdm
 
+
+logger = logging.getLogger(__name__)
+
 UPLOAD_CREDS = "/file/temporaryAccess"
 UPLOAD_CONFIRM = "/queue/confirmUpload"
 UPLOAD_CANCEL = "/file/cancelUpload"
@@ -45,12 +52,6 @@ class UploadCredentials(NamedTuple):
     bucket: str
 
 
-class FileUploadJob(NamedTuple):
-    mission_id: UUID
-    name: str
-    path: Path
-
-
 def _get_s3_endpoint() -> str:
     config = Config()
     endpoint = config.endpoint
@@ -69,9 +70,6 @@ def _confirm_file_upload(
         "md5": file_hash,
     }
     resp = client.post(UPLOAD_CONFIRM, json=data)
-
-    if 400 <= resp.status_code < 500:
-        raise CorruptedFile()
     resp.raise_for_status()
 
 
@@ -87,67 +85,47 @@ def _cancel_file_upload(
     return
 
 
-def _get_file_download(client: AuthenticatedClient, id: UUID) -> str:
-    """\
-    get the download url for a file by file id
-    """
-    resp = client.get(DOWNLOAD_URL, params={"uuid": str(id), "expires": True})
-
-    if 400 <= resp.status_code < 500:
-        raise AccessDeniedException(
-            f"Failed to download file: {resp.json()['message']}",
-            "Status Code: " + str(resp.status_code),
-        )
-
-    resp.raise_for_status()
+FILE_EXISTS_ERROR = "File already exists"
 
-    return resp.text
+# fields for upload credentials
+ACCESS_KEY_FIELD = "accessKey"
+SECRET_KEY_FIELD = "secretKey"
+SESSION_TOKEN_FIELD = "sessionToken"
+CREDENTIALS_FIELD = "accessCredentials"
+FILE_ID_FIELD = "fileUUID"
+BUCKET_FIELD = "bucket"
 
 
 def _get_upload_creditials(
-    client: AuthenticatedClient, internal_filenames: List[str], mission_id: UUID
-) -> Dict[str, UploadCredentials]:
-    if mission_id.version != 4:
-        raise ValueError("Mission ID must be a UUIDv4")
+    client: AuthenticatedClient, internal_filename: str, mission_id: UUID
+) -> Optional[UploadCredentials]:
     dct = {
-        "filenames": internal_filenames,
+        "filenames": [internal_filename],
         "missionUUID": str(mission_id),
     }
     resp = client.post(UPLOAD_CREDS, json=dct)
+    resp.raise_for_status()
 
-    if resp.status_code >= 400:
-        raise ValueError(
-            "Failed to get temporary credentials. Status Code: "
-            f"{resp.status_code}\n{resp.json()['message'][0]}"
-        )
-
-    data = resp.json()
-
-    ret = {}
-    for record in data:
-        if "error" in record:
-            # TODO: handle this better
-            continue
-
-        bucket = record["bucket"]
-        file_id = UUID(record["fileUUID"], version=4)
-        filename = record["fileName"]
+    data = resp.json()[0]
 
-        creds = record["accessCredentials"]
+    if data.get("error") == FILE_EXISTS_ERROR:
+        return None
 
-        access_key = creds["accessKey"]
-        secret_key = creds["secretKey"]
-        session_token = creds["sessionToken"]
+    bucket = data[BUCKET_FIELD]
+    file_id = UUID(data[FILE_ID_FIELD], version=4)
 
-        ret[filename] = UploadCredentials(
-            access_key=access_key,
-            secret_key=secret_key,
-            session_token=session_token,
-            file_id=file_id,
-            bucket=bucket,
-        )
+    creds = data[CREDENTIALS_FIELD]
+    access_key = creds[ACCESS_KEY_FIELD]
+    secret_key = creds[SECRET_KEY_FIELD]
+    session_token = creds[SESSION_TOKEN_FIELD]
 
-    return ret
+    return UploadCredentials(
+        access_key=access_key,
+        secret_key=secret_key,
+        session_token=session_token,
+        file_id=file_id,
+        bucket=bucket,
+    )
 
 
 def _s3_upload(
@@ -156,182 +134,333 @@ def _s3_upload(
     endpoint: str,
     credentials: UploadCredentials,
     pbar: tqdm,
-) -> bool:
+) -> None:
     # configure boto3
-    try:
-        config = botocore.config.Config(
-            retries={"max_attempts": S3_MAX_RETRIES},
-            read_timeout=S3_READ_TIMEOUT,
-        )
-        client = boto3.client(
-            "s3",
-            endpoint_url=endpoint,
-            aws_access_key_id=credentials.access_key,
-            aws_secret_access_key=credentials.secret_key,
-            aws_session_token=credentials.session_token,
-            config=config,
-        )
-        client.upload_file(
-            str(local_path),
-            credentials.bucket,
-            str(credentials.file_id),
-            Callback=pbar.update,
-        )
-    except Exception as e:
-        err = f"error uploading file: {local_path}: {type(e).__name__}"
-        pbar.write(raw_rich(Text(err, style="red")))
-        return False
-    return True
+    config = botocore.config.Config(
+        retries={"max_attempts": S3_MAX_RETRIES},
+        read_timeout=S3_READ_TIMEOUT,
+    )
+    client = boto3.client(
+        "s3",
+        endpoint_url=endpoint,
+        aws_access_key_id=credentials.access_key,
+        aws_secret_access_key=credentials.secret_key,
+        aws_session_token=credentials.session_token,
+        config=config,
+    )
+    client.upload_file(
+        str(local_path),
+        credentials.bucket,
+        str(credentials.file_id),
+        Callback=pbar.update,
+    )
 
 
-def _upload_file(
+class UploadState(Enum):
+    UPLOADED = 1
+    EXISTS = 2
+    CANCELED = 3
+
+
+# TODO: i dont want to handle errors at this level
+def upload_file(
     client: AuthenticatedClient,
-    job: FileUploadJob,
-    hide_progress: bool = False,
-    global_pbar: Optional[tqdm] = None,
-) -> Tuple[int, Path]:
+    *,
+    mission_id: UUID,
+    filename: str,
+    path: Path,
+    verbose: bool = False,
+) -> UploadState:
     """\
     returns bytes uploaded
     """
 
-    pbar = tqdm(
-        total=os.path.getsize(job.path),
+    total_size = path.stat().st_size
+    with tqdm(
+        total=total_size,
         unit="B",
         unit_scale=True,
-        desc=f"uploading {job.path.name}...",
+        desc=f"uploading {path}...",
         leave=False,
-        disable=hide_progress,
-    )
+        disable=not verbose,
+    ) as pbar:
+        endpoint = _get_s3_endpoint()
 
-    # get creditials for the upload
-    try:
-        # get upload credentials for a single file
-        access = _get_upload_creditials(
-            client, internal_filenames=[job.name], mission_id=job.mission_id
+        # get per file upload credentials
+        creds = _get_upload_creditials(
+            client, internal_filename=filename, mission_id=mission_id
         )
-        # upload file
-        creds = access[job.name]
-    except Exception as e:
-        pbar.write(f"unable to get upload credentials for file {job.path.name}: {e}")
-        pbar.close()
-        if global_pbar is not None:
-            global_pbar.update()
-        return (0, job.path)
-
-    # do the upload
-    endpoint = _get_s3_endpoint()
-    success = _s3_upload(job.path, endpoint=endpoint, credentials=creds, pbar=pbar)
+        if creds is None:
+            return UploadState.EXISTS
 
-    if not success:
-        try:
-            _cancel_file_upload(client, creds.file_id, job.mission_id)
-        except Exception as e:
-            msg = Text(f"failed to cancel upload: {type(e).__name__}", style="red")
-            pbar.write(raw_rich(msg))
-    else:
-        # tell backend that upload is complete
         try:
-            local_hash = b64_md5(job.path)
-            _confirm_file_upload(client, creds.file_id, local_hash)
-
-            if global_pbar is not None:
-                msg = Text(f"uploaded {job.path}", style="green")
-                global_pbar.write(raw_rich(msg))
-                global_pbar.update()
-
+            _s3_upload(path, endpoint=endpoint, credentials=creds, pbar=pbar)
         except Exception as e:
-            msg = Text(
-                f"error confirming upload {job.path}: {type(e).__name__}", style="red"
-            )
-            pbar.write(raw_rich(msg))
+            logger.error(format_traceback(e))
+            _cancel_file_upload(client, creds.file_id, mission_id)
+            return UploadState.CANCELED
 
-    pbar.close()
-    return (job.path.stat().st_size, job.path)
+        else:
+            _confirm_file_upload(client, creds.file_id, b64_md5(path))
+            return UploadState.UPLOADED
 
 
-def upload_files(
-    files_map: Dict[str, Path],
-    mission_id: UUID,
-    *,
-    verbose: bool = False,
-    n_workers: int = 2,
-) -> None:
-    futures = []
-
-    pbar = tqdm(
-        total=len(files_map),
-        unit="files",
-        desc="Uploading files",
-        disable=not verbose,
-    )
-
-    start = monotonic()
-    with ThreadPoolExecutor(max_workers=n_workers) as executor:
-        for name, path in files_map.items():
-            # client is not thread safe
-            client = AuthenticatedClient()
-            job = FileUploadJob(mission_id=mission_id, name=name, path=path)
-            future = executor.submit(
-                _upload_file,
-                client=client,
-                job=job,
-                hide_progress=not verbose,
-                global_pbar=pbar,
-            )
-            futures.append(future)
-
-    errors = []
-    total_size = 0
-    for f in futures:
-        try:
-            size, path = f.result()
-            size = size / 1024 / 1024  # convert to MB
-
-            if not verbose and size > 0:
-                print(path.absolte())
-
-            total_size += size
-        except Exception as e:
-            errors.append(e)
+def _get_file_download(client: AuthenticatedClient, id: UUID) -> str:
+    """\
+    get the download url for a file by file id
+    """
+    resp = client.get(DOWNLOAD_URL, params={"uuid": str(id), "expires": True})
 
-    pbar.close()
+    if 400 <= resp.status_code < 500:
+        raise AccessDenied(
+            f"Failed to download file: {resp.json()['message']}"
+            f"Status Code: {resp.status_code}",
+        )
 
-    time = monotonic() - start
-    print(f"upload took {time:.2f} seconds", file=sys.stderr)
-    print(f"total size: {int(total_size)} MB", file=sys.stderr)
-    print(f"average speed: {total_size / time:.2f} MB/s", file=sys.stderr)
+    resp.raise_for_status()
 
-    if errors:
-        raise UploadFailed(f"got unhandled errors: {errors} when uploading files")
+    return resp.text
 
 
-def _url_download(url: str, path: Path, size: int, overwrite: bool = False) -> None:
+def _url_download(
+    url: str, *, path: Path, size: int, overwrite: bool = False, verbose: bool = False
+) -> None:
     if path.exists() and not overwrite:
-        raise FileExistsError(f"File already exists: {path}")
+        raise FileExistsError(f"file already exists: {path}")
 
     with httpx.stream("GET", url) as response:
         with open(path, "wb") as f:
             with tqdm(
-                total=size, desc=f"Downloading {path.name}", unit="B", unit_scale=True
+                total=size,
+                desc=f"downloading {path.name}",
+                unit="B",
+                unit_scale=True,
+                leave=False,
+                disable=not verbose,
             ) as pbar:
                 for chunk in response.iter_bytes(chunk_size=DOWNLOAD_CHUNK_SIZE):
                     f.write(chunk)
                     pbar.update(len(chunk))
 
 
+class DownloadState(Enum):
+    DOWNLOADED_OK = 1
+    SKIPPED_OK = 2
+    DOWNLOADED_INVALID_HASH = 3
+    SKIPPED_INVALID_HASH = 4
+    SKIPPED_INVALID_REMOTE_STATE = 5
+
+
 def download_file(
     client: AuthenticatedClient,
-    file_id: UUID,
-    name: str,
-    dest: Path,
-    hash: str,
-    size: int,
-) -> None:
-    download_url = _get_file_download(client, file_id)
+    *,
+    file: File,
+    path: Path,
+    overwrite: bool = False,
+    verbose: bool = False,
+) -> DownloadState:
+    # skip files that are not ok on remote
+    if file.state != FileState.OK:
+        return DownloadState.SKIPPED_INVALID_REMOTE_STATE
+
+    # skip existing files depending on flags set
+    if path.exists():
+        local_hash = b64_md5(path)
+        if local_hash != file.hash and not overwrite and file.hash is not None:
+            return DownloadState.SKIPPED_INVALID_HASH
+
+        elif local_hash == file.hash:
+            return DownloadState.SKIPPED_OK
+
+        # this has to be here
+        if verbose:
+            tqdm.write(
+                styled_string(f"overwriting {path}, hash missmatch", style="yellow")
+            )
+
+    # request a download url
+    download_url = _get_file_download(client, file.id)
+
+    # create parent directories
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # download the file and check the hash
+    _url_download(
+        download_url, path=path, size=file.size, overwrite=overwrite, verbose=verbose
+    )
+    observed_hash = b64_md5(path)
+    if file.hash is not None and observed_hash != file.hash:
+        return DownloadState.DOWNLOADED_INVALID_HASH
+    return DownloadState.DOWNLOADED_OK
+
+
+UPLOAD_STATE_COLOR = {
+    UploadState.UPLOADED: "green",
+    UploadState.EXISTS: "yellow",
+    UploadState.CANCELED: "red",
+}
+
+
+def _upload_handler(
+    future: Future[UploadState], path: Path, *, verbose: bool = False
+) -> int:
+    try:
+        state = future.result()
+    except Exception as e:
+        logger.error(format_traceback(e))
+        if verbose:
+            tqdm.write(format_error(f"error uploading {path}", e))
+        else:
+            print(path.absolute(), file=sys.stderr)
+        return 0
+
+    if state == UploadState.UPLOADED:
+        msg = f"uploaded {path}"
+    elif state == UploadState.EXISTS:
+        msg = f"skipped {path} already uploaded"
+    else:
+        msg = f"canceled {path} upload"
+
+    if verbose:
+        tqdm.write(styled_string(msg, style=UPLOAD_STATE_COLOR[state]))
+    else:
+        stream = sys.stdout if state == UploadState.UPLOADED else sys.stderr
+        print(path.absolute(), file=stream)
 
-    file_path = dest / name
-    _url_download(download_url, file_path, size)
-    observed_hash = b64_md5(file_path)
+    return path.stat().st_size if state == UploadState.UPLOADED else 0
 
-    if observed_hash != hash:
-        raise CorruptedFile("file hash does not match")
+
+DOWNLOAD_STATE_COLOR = {
+    DownloadState.DOWNLOADED_OK: "green",
+    DownloadState.SKIPPED_OK: "green",
+    DownloadState.DOWNLOADED_INVALID_HASH: "red",
+    DownloadState.SKIPPED_INVALID_HASH: "yellow",
+    DownloadState.SKIPPED_INVALID_REMOTE_STATE: "purple",
+}
+
+
+def _download_handler(
+    future: Future[DownloadState], file: File, path: Path, *, verbose: bool = False
+) -> int:
+    try:
+        state = future.result()
+    except Exception as e:
+        logger.error(format_traceback(e))
+        if verbose:
+            tqdm.write(format_error(f"error uploading {path}", e))
+        else:
+            print(path.absolute(), file=sys.stderr)
+        return 0
+
+    if state == DownloadState.DOWNLOADED_OK:
+        msg = f"downloaded {path}"
+    elif state == DownloadState.DOWNLOADED_INVALID_HASH:
+        msg = f"downloaded {path} failed hash check"
+    elif state == DownloadState.SKIPPED_OK:
+        msg = f"skipped {path} already downloaded"
+    elif state == DownloadState.SKIPPED_INVALID_HASH:
+        msg = f"skipped {path} already downloaded, hash missmatch, cosider using `--overwrite`"
+    else:
+        msg = f"skipped {path} remote file has invalid state"
+
+    if verbose:
+        tqdm.write(styled_string(msg, style=DOWNLOAD_STATE_COLOR[state]))
+    else:
+        stream = (
+            sys.stdout
+            if state in (DownloadState.DOWNLOADED_OK, DownloadState.SKIPPED_OK)
+            else sys.stderr
+        )
+        print(path.absolute(), file=stream)
+
+    # number of bytes downloaded
+    return file.size if state == DownloadState.DOWNLOADED_OK else 0
+
+
+def upload_files(
+    client: AuthenticatedClient,
+    files_map: Dict[str, Path],
+    mission_id: UUID,
+    *,
+    verbose: bool = False,
+    n_workers: int = 2,
+) -> None:
+    with tqdm(
+        total=len(files_map),
+        unit="files",
+        desc="uploading files",
+        disable=not verbose,
+        leave=False,
+    ) as pbar:
+        start = monotonic()
+        futures: Dict[Future[UploadState], Path] = {}
+        with ThreadPoolExecutor(max_workers=n_workers) as executor:
+            for name, path in files_map.items():
+                future = executor.submit(
+                    upload_file,
+                    client=client,
+                    mission_id=mission_id,
+                    filename=name,
+                    path=path,
+                    verbose=verbose,
+                )
+                futures[future] = path
+
+            total_size = 0
+            for future in as_completed(futures):
+                size = _upload_handler(future, futures[future], verbose=verbose)
+                total_size += size / 1024 / 1024
+
+                pbar.update()
+            pbar.refresh()
+
+        t = monotonic() - start
+        c = Console(file=sys.stderr)
+        c.print(f"upload took {t:.2f} seconds")
+        c.print(f"total size: {int(total_size)} MB")
+        c.print(f"average speed: {total_size / t:.2f} MB/s")
+
+
+def download_files(
+    client: AuthenticatedClient,
+    files: Dict[Path, File],
+    *,
+    verbose: bool = False,
+    overwrite: bool = False,
+    n_workers: int = 2,
+) -> None:
+    with tqdm(
+        total=len(files),
+        unit="files",
+        desc="downloading files",
+        disable=not verbose,
+        leave=False,
+    ) as pbar:
+
+        start = monotonic()
+        futures: Dict[Future[DownloadState], Tuple[File, Path]] = {}
+        with ThreadPoolExecutor(max_workers=n_workers) as executor:
+            for path, file in files.items():
+                future = executor.submit(
+                    download_file,
+                    client=client,
+                    file=file,
+                    path=path,
+                    overwrite=overwrite,
+                    verbose=verbose,
+                )
+                futures[future] = (file, path)
+
+            total_size = 0
+            for future in as_completed(futures):
+                file, path = futures[future]
+                size = _download_handler(future, file, path, verbose=verbose)
+                total_size += size / 1024 / 1024  # MB
+                pbar.update()
+            pbar.refresh()
+
+        time = monotonic() - start
+        c = Console(file=sys.stderr)
+        c.print(f"download took {time:.2f} seconds")
+        c.print(f"total size: {int(total_size)} MB")
+        c.print(f"average speed: {total_size  / time:.2f} MB/s")