kleinkram 0.13.1__tar.gz → 0.13.2__tar.gz

{kleinkram-0.13.1 → kleinkram-0.13.2}/PKG-INFO

@@ -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 → kleinkram-0.13.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "kleinkram"
-version = "0.13.1"
+version = "0.13.2"
 description = "A CLI for the ETH project kleinkram"
 readme = "README.md"
 requires-python = ">=3.7"

{kleinkram-0.13.1 → kleinkram-0.13.2}/requirements.txt

@@ -5,4 +5,6 @@ typing_extensions~=4.12.2
 rich~=13.7.1
 tqdm~=4.66.4
 black~=24.8.0
-requests~=2.32.3
+requests~=2.32.3
+boto3~=1.35.9
+botocore~=1.35.9

kleinkram-0.13.2/src/kleinkram/helper.py

@@ -0,0 +1,169 @@
+import glob
+import os
+import queue
+import threading
+from functools import partial
+from typing import Dict
+import boto3
+
+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)
+
+    normalized_path = os.path.normpath(expanded_path)
+
+    if "**" in normalized_path:
+        file_list = glob.glob(normalized_path, recursive=True)
+    else:
+        file_list = glob.glob(normalized_path)
+
+    return file_list
+
+
+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 = []
+    transferCallback = TransferCallback()
+
+    for i in range(nrThreads):
+        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, s3: BaseClient, transferCallback: TransferCallback):
+    while True:
+        try:
+            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:
+                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
+        except Exception as e:
+            print(f"Error uploading {filename}: {e}")
+            _queue.task_done()
+
+
+if __name__ == "__main__":
+    res = expand_and_match(
+        "~/Downloads/dodo_mission_2024_02_08-20240408T074313Z-003/**.bag"
+    )
+    print(res)

{kleinkram-0.13.1 → kleinkram-0.13.2}/src/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/src/kleinkram/helper.py

@@ -1,74 +0,0 @@
-import glob
-import os
-import queue
-import threading
-from typing import Dict
-
-import httpx
-import tqdm
-from rich import print
-
-from kleinkram.api_client import AuthenticatedClient
-
-
-def expand_and_match(path_pattern):
-    expanded_path = os.path.expanduser(path_pattern)
-    expanded_path = os.path.expandvars(expanded_path)
-
-    normalized_path = os.path.normpath(expanded_path)
-
-    if "**" in normalized_path:
-        file_list = glob.glob(normalized_path, recursive=True)
-    else:
-        file_list = glob.glob(normalized_path)
-
-    return file_list
-
-
-def uploadFiles(files: Dict[str, str], paths: Dict[str, str], nrThreads: int):
-    _queue = queue.Queue()
-    for file in files.items():
-        _queue.put(file)
-    threads = []
-    pbar = tqdm.tqdm(total=len(files.items()) * 100)
-    for i in range(nrThreads):
-        thread = threading.Thread(target=uploadFile, args=(_queue, paths, pbar))
-        thread.start()
-        threads.append(thread)
-    for thread in threads:
-        thread.join()
-
-
-def uploadFile(_queue: queue.Queue, paths: Dict[str, str], pbar: tqdm):
-    while True:
-        try:
-            filename, info = _queue.get(timeout=3)
-            url = info["url"]
-            uuid = info["uuid"]
-            filepath = paths[filename]
-            headers = {"Content-Type": "application/octet-stream"}
-            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}"
-                        )
-            _queue.task_done()
-        except queue.Empty:
-            break
-        except Exception as e:
-            print(f"Error uploading {filename}: {e}")
-            _queue.task_done()
-
-
-if __name__ == "__main__":
-    res = expand_and_match(
-        "~/Downloads/dodo_mission_2024_02_08-20240408T074313Z-003/**.bag"
-    )
-    print(res)