kleinkram 0.37.0.dev20241113182530__tar.gz → 0.37.0.dev20241118113347__tar.gz

{kleinkram-0.37.0.dev20241113182530 → kleinkram-0.37.0.dev20241118113347}/LICENSE

@@ -671,4 +671,4 @@ into proprietary programs.  If your program is a subroutine library, you
 may consider it more useful to permit linking proprietary applications with
 the library.  If this is what you want to do, use the GNU Lesser General
 Public License instead of this License.  But first, please read
-<https://www.gnu.org/licenses/why-not-lgpl.html>.
+<https://www.gnu.org/licenses/why-not-lgpl.html>.

kleinkram-0.37.0.dev20241118113347/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.1
+Name: kleinkram
+Version: 0.37.0.dev20241118113347
+Summary: give me your bags
+Author: Cyrill Püntener, Dominique Garmier, Johann Schwabe
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: boto3
+Requires-Dist: botocore
+Requires-Dist: httpx
+Requires-Dist: pyyaml
+Requires-Dist: rich
+Requires-Dist: tqdm
+Requires-Dist: typer
+
+# Kleinkram: CLI
+
+Install the package
+
+```bash
+pip install kleinkram
+```
+
+Run the CLI
+
+```bash
+klein
+```
+
+## Usage
+
+Here are some basic examples of how to use the CLI.
+
+### Listing Files
+
+To list all files in a mission:
+
+```bash
+klein list -p project-name -m mission-name
+```
+
+### Uploading Files
+
+To upload all `*.bag` files in the current directory to a mission:
+
+```bash
+klein upload -p project-name -m mission-name *.bag
+```
+
+If you would like to create a new mission on upload use the `--create` flag.
+
+### Downloading Files
+
+To download all files from a mission and save them `out`:
+
+```bash
+klein download -p project-name -m mission-name --dest out
+```
+
+You can additionally specify filenames or ids if you only want to download specific files.
+
+Instead of downloading files from a specified mission you can download arbitrary files by specifying their ids:
+
+```bash
+klein download --dest out *id1* *id2* *id3*
+```
+
+For more information consult the [documentation](https://docs.datasets.leggedrobotics.com/usage/cli/cli-getting-started.html).
+
+## Development
+
+Clone the repo
+
+```bash
+git clone git@github.com:leggedrobotics/kleinkram.git
+cd kleinkram/cli
+```
+
+Setup the environment
+
+```bash
+virtualenv -ppython3.8 .venv
+source .venv/bin/activate
+pip install -e . -r requirements.txt
+```
+
+Install `pre-commit` hooks
+
+```bash
+pre-commit install
+```
+
+Run the CLI
+
+```bash
+klein --help
+```
+
+### Run Tests
+```bash
+pytest .
+```
+
+You can also look in `scripts` for some scripts that might be useful for testing.

kleinkram-0.37.0.dev20241118113347/README.md

@@ -0,0 +1,89 @@
+# Kleinkram: CLI
+
+Install the package
+
+```bash
+pip install kleinkram
+```
+
+Run the CLI
+
+```bash
+klein
+```
+
+## Usage
+
+Here are some basic examples of how to use the CLI.
+
+### Listing Files
+
+To list all files in a mission:
+
+```bash
+klein list -p project-name -m mission-name
+```
+
+### Uploading Files
+
+To upload all `*.bag` files in the current directory to a mission:
+
+```bash
+klein upload -p project-name -m mission-name *.bag
+```
+
+If you would like to create a new mission on upload use the `--create` flag.
+
+### Downloading Files
+
+To download all files from a mission and save them `out`:
+
+```bash
+klein download -p project-name -m mission-name --dest out
+```
+
+You can additionally specify filenames or ids if you only want to download specific files.
+
+Instead of downloading files from a specified mission you can download arbitrary files by specifying their ids:
+
+```bash
+klein download --dest out *id1* *id2* *id3*
+```
+
+For more information consult the [documentation](https://docs.datasets.leggedrobotics.com/usage/cli/cli-getting-started.html).
+
+## Development
+
+Clone the repo
+
+```bash
+git clone git@github.com:leggedrobotics/kleinkram.git
+cd kleinkram/cli
+```
+
+Setup the environment
+
+```bash
+virtualenv -ppython3.8 .venv
+source .venv/bin/activate
+pip install -e . -r requirements.txt
+```
+
+Install `pre-commit` hooks
+
+```bash
+pre-commit install
+```
+
+Run the CLI
+
+```bash
+klein --help
+```
+
+### Run Tests
+```bash
+pytest .
+```
+
+You can also look in `scripts` for some scripts that might be useful for testing.

kleinkram-0.37.0.dev20241118113347/kleinkram/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from kleinkram._version import __version__
+
+
+__all__ = ["__version__"]

kleinkram-0.37.0.dev20241118113347/kleinkram/__main__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from kleinkram.main import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

kleinkram-0.37.0.dev20241118113347/kleinkram/_version.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from importlib.metadata import version
+
+__version__ = version("kleinkram")
+__local__ = False

kleinkram-0.37.0.dev20241118113347/kleinkram/api/client.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import httpx
+from kleinkram.auth import Config
+from kleinkram.config import Credentials
+from kleinkram.errors import LOGIN_MESSAGE
+from kleinkram.errors import NotAuthenticatedException
+
+
+COOKIE_AUTH_TOKEN = "authtoken"
+COOKIE_REFRESH_TOKEN = "refreshtoken"
+COOKIE_CLI_KEY = "clikey"
+
+
+class AuthenticatedClient(httpx.Client):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.config = Config()
+
+        if self.config.has_cli_key:
+            assert self.config.cli_key, "unreachable"
+            self.cookies.set(COOKIE_CLI_KEY, self.config.cli_key)
+
+        elif self.config.has_refresh_token:
+            assert self.config.auth_token is not None, "unreachable"
+            self.cookies.set(COOKIE_AUTH_TOKEN, self.config.auth_token)
+        else:
+            raise NotAuthenticatedException(self.config.endpoint)
+
+    def _refresh_token(self) -> None:
+        if self.config.has_cli_key:
+            raise RuntimeError
+
+        refresh_token = self.config.refresh_token
+        if not refresh_token:
+            raise RuntimeError
+
+        self.cookies.set(COOKIE_REFRESH_TOKEN, refresh_token)
+
+        response = self.post(
+            "/auth/refresh-token",
+        )
+        response.raise_for_status()
+
+        new_access_token = response.cookies[COOKIE_AUTH_TOKEN]
+        creds = Credentials(auth_token=new_access_token, refresh_token=refresh_token)
+
+        self.config.save_credentials(creds)
+        self.cookies.set(COOKIE_AUTH_TOKEN, new_access_token)
+
+    def request(self, method, url, *args, **kwargs):
+        full_url = f"{self.config.endpoint}{url}"
+        response = super().request(method, full_url, *args, **kwargs)
+
+        if (url == "/auth/refresh-token") and response.status_code == 401:
+            raise RuntimeError(LOGIN_MESSAGE)
+
+        if response.status_code == 401:
+            try:
+                self._refresh_token()
+            except Exception:
+                raise RuntimeError(LOGIN_MESSAGE)
+            return super().request(method, full_url, *args, **kwargs)
+        else:
+            return response

kleinkram-0.37.0.dev20241118113347/kleinkram/api/file_transfer.py

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+import os
+import sys
+from concurrent.futures import ThreadPoolExecutor
+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
+from uuid import UUID
+
+import boto3.s3.transfer
+import botocore.config
+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.utils import b64_md5
+from kleinkram.utils import raw_rich
+from rich.text import Text
+from tqdm import tqdm
+
+UPLOAD_CREDS = "/file/temporaryAccess"
+UPLOAD_CONFIRM = "/queue/confirmUpload"
+UPLOAD_CANCEL = "/file/cancelUpload"
+
+DOWNLOAD_CHUNK_SIZE = 1024 * 1024 * 16
+DOWNLOAD_URL = "/file/download"
+
+S3_MAX_RETRIES = 60  # same as frontend
+S3_READ_TIMEOUT = 60 * 5  # 5 minutes
+
+
+class UploadCredentials(NamedTuple):
+    access_key: str
+    secret_key: str
+    session_token: str
+    file_id: UUID
+    bucket: str
+
+
+class FileUploadJob(NamedTuple):
+    mission_id: UUID
+    name: str
+    path: Path
+
+
+def _get_s3_endpoint() -> str:
+    config = Config()
+    endpoint = config.endpoint
+
+    if "localhost" in endpoint:
+        return LOCAL_S3
+    else:
+        return endpoint.replace("api", "minio")
+
+
+def _confirm_file_upload(
+    client: AuthenticatedClient, file_id: UUID, file_hash: str
+) -> None:
+    data = {
+        "uuid": str(file_id),
+        "md5": file_hash,
+    }
+    resp = client.post(UPLOAD_CONFIRM, json=data)
+
+    if 400 <= resp.status_code < 500:
+        raise CorruptedFile()
+    resp.raise_for_status()
+
+
+def _cancel_file_upload(
+    client: AuthenticatedClient, file_id: UUID, mission_id: UUID
+) -> None:
+    data = {
+        "uuid": [str(file_id)],
+        "missionUUID": str(mission_id),
+    }
+    resp = client.post(UPLOAD_CANCEL, json=data)
+    resp.raise_for_status()
+    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()
+
+    return resp.text
+
+
+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")
+    dct = {
+        "filenames": internal_filenames,
+        "missionUUID": str(mission_id),
+    }
+    resp = client.post(UPLOAD_CREDS, json=dct)
+
+    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"]
+
+        creds = record["accessCredentials"]
+
+        access_key = creds["accessKey"]
+        secret_key = creds["secretKey"]
+        session_token = creds["sessionToken"]
+
+        ret[filename] = UploadCredentials(
+            access_key=access_key,
+            secret_key=secret_key,
+            session_token=session_token,
+            file_id=file_id,
+            bucket=bucket,
+        )
+
+    return ret
+
+
+def _s3_upload(
+    local_path: Path,
+    *,
+    endpoint: str,
+    credentials: UploadCredentials,
+    pbar: tqdm,
+) -> bool:
+    # 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
+
+
+def _upload_file(
+    client: AuthenticatedClient,
+    job: FileUploadJob,
+    hide_progress: bool = False,
+    global_pbar: Optional[tqdm] = None,
+) -> Tuple[int, Path]:
+    """\
+    returns bytes uploaded
+    """
+
+    pbar = tqdm(
+        total=os.path.getsize(job.path),
+        unit="B",
+        unit_scale=True,
+        desc=f"uploading {job.path.name}...",
+        leave=False,
+        disable=hide_progress,
+    )
+
+    # 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
+        )
+        # 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 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()
+
+        except Exception as e:
+            msg = Text(
+                f"error confirming upload {job.path}: {type(e).__name__}", style="red"
+            )
+            pbar.write(raw_rich(msg))
+
+    pbar.close()
+    return (job.path.stat().st_size, job.path)
+
+
+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)
+
+    pbar.close()
+
+    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)
+
+    if errors:
+        raise UploadFailed(f"got unhandled errors: {errors} when uploading files")
+
+
+def _url_download(url: str, path: Path, size: int, overwrite: bool = False) -> None:
+    if path.exists() and not overwrite:
+        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
+            ) as pbar:
+                for chunk in response.iter_bytes(chunk_size=DOWNLOAD_CHUNK_SIZE):
+                    f.write(chunk)
+                    pbar.update(len(chunk))
+
+
+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_path = dest / name
+    _url_download(download_url, file_path, size)
+    observed_hash = b64_md5(file_path)
+
+    if observed_hash != hash:
+        raise CorruptedFile("file hash does not match")