kleinkram 0.13.1__py3-none-any.whl → 0.13.2__py3-none-any.whl

kleinkram/helper.py

@@ -2,15 +2,91 @@ import glob
 import os
 import queue
 import threading
+from functools import partial
 from typing import Dict
+import boto3
 
-import httpx
 import tqdm
+from boto3.s3.transfer import TransferConfig
+from botocore.client import BaseClient
 from rich import print
 
 from kleinkram.api_client import AuthenticatedClient
 
 
+class TransferCallback:
+    """
+    Handle callbacks from the transfer manager.
+
+    The transfer manager periodically calls the __call__ method throughout
+    the upload process so that it can take action, such as displaying progress
+    to the user and collecting data about the transfer.
+    """
+
+    def __init__(self):
+        """
+        Initialize the TransferCallback.
+
+        This initializes an empty dictionary to hold progress bars for each file.
+        """
+        self._lock = threading.Lock()
+        self.file_progress = {}
+
+    def add_file(self, file_id, target_size):
+        """
+        Add a new file to track.
+
+        :param file_id: A unique identifier for the file (e.g., file name or ID).
+        :param target_size: The total size of the file being transferred.
+        """
+        with self._lock:
+            tqdm_instance = tqdm.tqdm(
+                total=target_size,
+                unit="B",
+                unit_scale=True,
+                desc=f"Uploading {file_id}",
+            )
+            self.file_progress[file_id] = {
+                "tqdm": tqdm_instance,
+                "total_transferred": 0,
+            }
+
+    def __call__(self, file_id, bytes_transferred):
+        """
+        The callback method that is called by the transfer manager.
+
+        Display progress during file transfer and collect per-thread transfer
+        data. This method can be called by multiple threads, so shared instance
+        data is protected by a thread lock.
+
+        :param file_id: The identifier of the file being transferred.
+        :param bytes_transferred: The number of bytes transferred in this call.
+        """
+        with self._lock:
+            if file_id in self.file_progress:
+                progress = self.file_progress[file_id]
+                progress["total_transferred"] += bytes_transferred
+
+                # Update tqdm progress bar
+                progress["tqdm"].update(bytes_transferred)
+
+    def close(self):
+        """Close all tqdm progress bars."""
+        with self._lock:
+            for progress in self.file_progress.values():
+                progress["tqdm"].close()
+
+
+def create_transfer_callback(callback_instance, file_id):
+    """
+    Factory function to create a partial function for TransferCallback.
+    :param callback_instance: Instance of TransferCallback.
+    :param file_id: The unique identifier for the file.
+    :return: A callable that can be passed as a callback to boto3's upload_file method.
+    """
+    return partial(callback_instance.__call__, file_id)
+
+
 def expand_and_match(path_pattern):
     expanded_path = os.path.expanduser(path_pattern)
     expanded_path = os.path.expandvars(expanded_path)
@@ -25,40 +101,59 @@ def expand_and_match(path_pattern):
     return file_list
 
 
-def uploadFiles(files: Dict[str, str], paths: Dict[str, str], nrThreads: int):
+def uploadFiles(files: Dict[str, str], credentials: Dict[str, str], nrThreads: int):
+    client = AuthenticatedClient()
+
+    session = boto3.Session(
+        aws_access_key_id=credentials["accessKey"],
+        aws_secret_access_key=credentials["secretKey"],
+        aws_session_token=credentials["sessionToken"],
+    )
+    api_endpoint = client.tokenfile.endpoint
+    if api_endpoint == "http://localhost:3000":
+        minio_endpoint = "http://localhost:9000"
+    else:
+        minio_endpoint = api_endpoint.replace("api", "minio")
+    s3 = session.resource("s3", endpoint_url=minio_endpoint)
+
     _queue = queue.Queue()
     for file in files.items():
         _queue.put(file)
     threads = []
-    pbar = tqdm.tqdm(total=len(files.items()) * 100)
+    transferCallback = TransferCallback()
+
     for i in range(nrThreads):
-        thread = threading.Thread(target=uploadFile, args=(_queue, paths, pbar))
+        thread = threading.Thread(
+            target=uploadFile, args=(_queue, s3, transferCallback)
+        )
         thread.start()
         threads.append(thread)
     for thread in threads:
         thread.join()
 
 
-def uploadFile(_queue: queue.Queue, paths: Dict[str, str], pbar: tqdm):
+def uploadFile(_queue: queue.Queue, s3: BaseClient, transferCallback: TransferCallback):
     while True:
         try:
-            filename, info = _queue.get(timeout=3)
-            url = info["url"]
-            uuid = info["uuid"]
-            filepath = paths[filename]
-            headers = {"Content-Type": "application/octet-stream"}
+            filename, _file = _queue.get(timeout=3)
+            queueUUID = _file["queueUUID"]
+            filepath = _file["filepath"]
+            bucket = _file["bucket"]
+            target_location = _file["location"]
+            config = TransferConfig(
+                multipart_chunksize=10 * 1024 * 1024, max_concurrency=5
+            )
             with open(filepath, "rb") as f:
-                with httpx.Client() as cli:
-                    # Using PUT method directly for the upload
-                    response = cli.put(url, content=f, headers=headers)
-                    if response.status_code == 200:
-                        pbar.update(100)  # Update progress for each file
-                        client = AuthenticatedClient()
-                        client.post("/queue/confirmUpload", json={"uuid": uuid})
-                    else:
-                        print(
-                            f"Failed to upload {filename}. HTTP status: {response.status_code}"
-                        )
+                size = os.path.getsize(filepath)
+                transferCallback.add_file(filename, size)
+                callback_function = create_transfer_callback(transferCallback, filename)
+                s3.Bucket(bucket).upload_file(
+                    filepath, target_location, Config=config, Callback=callback_function
+                )
+
+                client = AuthenticatedClient()
+                res = client.post("/queue/confirmUpload", json={"uuid": queueUUID})
+                res.raise_for_status()
             _queue.task_done()
         except queue.Empty:
             break

kleinkram/main.py

@@ -179,23 +179,25 @@ def upload(
         )
         new_mission.raise_for_status()
         new_mission_data = new_mission.json()
-        print(f"Created mission: {new_mission_data['name']}")
 
-        get_presigned_url = "/queue/createPreSignedURLS"
+        get_temporary_credentials = "/file/temporaryAccess"
 
         response_2 = client.post(
-            get_presigned_url,
+            get_temporary_credentials,
             json={"filenames": filenames, "missionUUID": new_mission_data["uuid"]},
         )
         response_2.raise_for_status()
-        presigned_urls = response_2.json()
-        for file in filenames:
-            if not file in presigned_urls.keys():
+        temp_credentials = response_2.json()
+        credential = temp_credentials["credentials"]
+        confirmed_files = temp_credentials["files"]
+        for _file in filenames:
+            if not _file in confirmed_files.keys():
                 raise Exception(
                     "Could not upload File '" + file + "'. Is the filename unique? "
                 )
-        if len(presigned_urls) > 0:
-            uploadFiles(presigned_urls, filepaths, 4)
+            confirmed_files[_file]["filepath"] = filepaths[_file]
+        if len(confirmed_files.keys()) > 0:
+            uploadFiles(confirmed_files, credential, 4)
 
     except httpx.HTTPError as e:
         print(e)

{kleinkram-0.13.1.dist-info → kleinkram-0.13.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kleinkram
-Version: 0.13.1
+Version: 0.13.2
 Summary: A CLI for the ETH project kleinkram
 Project-URL: Homepage, https://github.com/leggedrobotics/kleinkram
 Project-URL: Issues, https://github.com/leggedrobotics/kleinkram/issues

{kleinkram-0.13.1.dist-info → kleinkram-0.13.2.dist-info}/RECORD

@@ -2,8 +2,8 @@ kleinkram/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kleinkram/api_client.py,sha256=1GPsM-XFbPYEKP7RfWmzMTwxRqnVh4wtHVuW25KT8kA,2264
 kleinkram/consts.py,sha256=pm_6OuQcO-tYcRhwauTtyRRsuYY0y0yb6EGuIl49LnI,50
 kleinkram/error_handling.py,sha256=Nm3tUtc4blQylJLNChSLgkDLFDHg3Xza0CghRrd_SYE,4015
-kleinkram/helper.py,sha256=9YUuCH0pfj7gK88XRtR0SD-cfdcbR-4g0DdAcUfwdm4,2334
-kleinkram/main.py,sha256=LIhlsepSWm_cl_VN0TuvPYtdKX_2Mz7d3Q-piebj3O8,7983
+kleinkram/helper.py,sha256=PNx6h7eTFZ10Ul7BqK2PTMXlKjHMKkQtANsceNZmAe8,5579
+kleinkram/main.py,sha256=ofSs0hFmi05A5gwa5t_ioVS6-aI0h27J8I1ENIxkZKk,8118
 kleinkram/auth/auth.py,sha256=bROptCsE0r5D416_7l1lfw52IX_mSVEnjaKiU2_b1Ms,4980
 kleinkram/endpoint/endpoint.py,sha256=KAYQgK8J6Et8V_ho2wBUIAdg25Mps6l_glCfC1vvC2g,1372
 kleinkram/file/file.py,sha256=HsKzg8xn-AHJJIF11H6c0Gy61cRX3MyTMYdcyypUxWo,3195
@@ -13,8 +13,8 @@ kleinkram/queue/queue.py,sha256=MaLBjAu8asi9BkPvbbT-5AobCcpy3ex5rxM1kHpRINA,181
 kleinkram/tag/tag.py,sha256=TZHh1GIV4LhwhqEGKdTJkM5mCIoLXAopBgWdLX1AEz8,1819
 kleinkram/topic/topic.py,sha256=qit-DECI19QocXEkkkBC0FAlzpZnotAtX8mylA9zqAA,1630
 kleinkram/user/user.py,sha256=i_QfsctjhImvKKjuDPfOIyDr322SXgV-KxJo-a7qNZw,1368
-kleinkram-0.13.1.dist-info/METADATA,sha256=hSkr3cBsUq6jdaTYGbvDZDOcbc1eqGMMkRcL0P5irNQ,756
-kleinkram-0.13.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-kleinkram-0.13.1.dist-info/entry_points.txt,sha256=RHXtRzcreVHImatgjhQwZQ6GdJThElYjHEWcR1BPXUI,45
-kleinkram-0.13.1.dist-info/licenses/LICENSE,sha256=ixuiBLtpoK3iv89l7ylKkg9rs2GzF9ukPH7ynZYzK5s,35148
-kleinkram-0.13.1.dist-info/RECORD,,
+kleinkram-0.13.2.dist-info/METADATA,sha256=ba5y3Udzl4-bZyf3MMu-qyCzRgI1wlJU2x5npP17TFc,756
+kleinkram-0.13.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+kleinkram-0.13.2.dist-info/entry_points.txt,sha256=RHXtRzcreVHImatgjhQwZQ6GdJThElYjHEWcR1BPXUI,45
+kleinkram-0.13.2.dist-info/licenses/LICENSE,sha256=ixuiBLtpoK3iv89l7ylKkg9rs2GzF9ukPH7ynZYzK5s,35148
+kleinkram-0.13.2.dist-info/RECORD,,