kleinkram 0.38.1.dev20241120100707__tar.gz → 0.38.1.dev20241125112529__tar.gz

{kleinkram-0.38.1.dev20241120100707 → kleinkram-0.38.1.dev20241125112529}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kleinkram
-Version: 0.38.1.dev20241120100707
+Version: 0.38.1.dev20241125112529
 Summary: give me your bags
 Author: Cyrill Püntener, Dominique Garmier, Johann Schwabe
 Classifier: Programming Language :: Python :: 3
@@ -109,5 +109,7 @@ klein --help
 ```bash
 pytest .
 ```
-
-You can also look in `scripts` for some scripts that might be useful for testing.
+or if you want to skip slow tests...
+```bash
+pytest -m "not slow" .
+```

{kleinkram-0.38.1.dev20241120100707 → kleinkram-0.38.1.dev20241125112529}/README.md

@@ -85,5 +85,7 @@ klein --help
 ```bash
 pytest .
 ```
-
-You can also look in `scripts` for some scripts that might be useful for testing.
+or if you want to skip slow tests...
+```bash
+pytest -m "not slow" .
+```

{kleinkram-0.38.1.dev20241120100707 → kleinkram-0.38.1.dev20241125112529}/kleinkram/api/client.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import logging
+from threading import Lock
 from typing import Any
 
 import httpx
 from kleinkram.auth import Config
 from kleinkram.config import Credentials
-from kleinkram.errors import LOGIN_MESSAGE
-from kleinkram.errors import NotAuthenticatedException
+from kleinkram.errors import NotAuthenticated
 
 logger = logging.getLogger(__name__)
 
@@ -21,31 +21,35 @@ class NotLoggedInException(Exception): ...
 
 
 class AuthenticatedClient(httpx.Client):
+    _config: Config
+    _config_lock: Lock
+
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        self.config = Config()
 
-        if self.config.has_cli_key:
-            assert self.config.cli_key, "unreachable"
+        self._config = Config()
+        self._config_lock = Lock()
+
+        if self._config.has_cli_key:
+            assert self._config.cli_key, "unreachable"
             logger.info("using cli key...")
-            self.cookies.set(COOKIE_CLI_KEY, self.config.cli_key)
+            self.cookies.set(COOKIE_CLI_KEY, self._config.cli_key)
 
-        elif self.config.has_refresh_token:
+        elif self._config.has_refresh_token:
             logger.info("using refresh token...")
-            assert self.config.auth_token is not None, "unreachable"
-            self.cookies.set(COOKIE_AUTH_TOKEN, self.config.auth_token)
+            assert self._config.auth_token is not None, "unreachable"
+            self.cookies.set(COOKIE_AUTH_TOKEN, self._config.auth_token)
         else:
             logger.info("not authenticated...")
-            raise NotAuthenticatedException(self.config.endpoint)
+            raise NotAuthenticated
 
     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
+        if self._config.has_cli_key:
+            raise RuntimeError("cannot refresh token when using cli key auth")
 
+        refresh_token = self._config.refresh_token
+        if refresh_token is None:
+            raise RuntimeError("no refresh token found")
         self.cookies.set(COOKIE_REFRESH_TOKEN, refresh_token)
 
         logger.info("refreshing token...")
@@ -53,30 +57,35 @@ class AuthenticatedClient(httpx.Client):
             "/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)
 
         logger.info("saving new tokens...")
-        self.config.save_credentials(creds)
+
+        with self._config_lock:
+            self._config.save_credentials(creds)
+
         self.cookies.set(COOKIE_AUTH_TOKEN, new_access_token)
 
     def request(
         self, method: str, url: str | httpx.URL, *args: Any, **kwargs: Any
     ) -> httpx.Response:
         if isinstance(url, httpx.URL):
-            raise ValueError("url must be a slug")
+            raise NotImplementedError(f"`httpx.URL` is not supported {url!r}")
+        if not url.startswith("/"):
+            url = f"/{url}"
 
         # try to do a request
-        full_url = f"{self.config.endpoint}{url}"
+        full_url = f"{self._config.endpoint}{url}"
         logger.info(f"requesting {method} {full_url}")
         response = super().request(method, full_url, *args, **kwargs)
+
         logger.info(f"got response {response}")
 
         # if the requesting a refresh token fails, we are not logged in
         if (url == "/auth/refresh-token") and response.status_code == 401:
             logger.info("got 401, not logged in...")
-            raise NotLoggedInException(LOGIN_MESSAGE)
+            raise NotAuthenticated
 
         # otherwise we try to refresh the token
         if response.status_code == 401:
@@ -84,12 +93,11 @@ class AuthenticatedClient(httpx.Client):
             try:
                 self._refresh_token()
             except Exception:
-                raise NotLoggedInException(LOGIN_MESSAGE)
+                raise NotAuthenticated
 
             logger.info(f"retrying request {method} {full_url}")
             resp = super().request(method, full_url, *args, **kwargs)
             logger.info(f"got response {resp}")
             return resp
-
         else:
             return response

kleinkram-0.38.1.dev20241125112529/kleinkram/api/file_transfer.py

@@ -0,0 +1,466 @@
+from __future__ import annotations
+
+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 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 AccessDenied
+from kleinkram.models import File
+from kleinkram.models import FileState
+from kleinkram.utils import b64_md5
+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"
+
+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
+
+
+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)
+    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
+
+
+FILE_EXISTS_ERROR = "File already exists"
+
+# 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_filename: str, mission_id: UUID
+) -> Optional[UploadCredentials]:
+    dct = {
+        "filenames": [internal_filename],
+        "missionUUID": str(mission_id),
+    }
+    resp = client.post(UPLOAD_CREDS, json=dct)
+    resp.raise_for_status()
+
+    data = resp.json()[0]
+
+    if data.get("error") == FILE_EXISTS_ERROR:
+        return None
+
+    bucket = data[BUCKET_FIELD]
+    file_id = UUID(data[FILE_ID_FIELD], version=4)
+
+    creds = data[CREDENTIALS_FIELD]
+    access_key = creds[ACCESS_KEY_FIELD]
+    secret_key = creds[SECRET_KEY_FIELD]
+    session_token = creds[SESSION_TOKEN_FIELD]
+
+    return UploadCredentials(
+        access_key=access_key,
+        secret_key=secret_key,
+        session_token=session_token,
+        file_id=file_id,
+        bucket=bucket,
+    )
+
+
+def _s3_upload(
+    local_path: Path,
+    *,
+    endpoint: str,
+    credentials: UploadCredentials,
+    pbar: tqdm,
+) -> None:
+    # configure boto3
+    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,
+    )
+
+
+class UploadState(Enum):
+    UPLOADED = 1
+    EXISTS = 2
+    CANCELED = 3
+
+
+# TODO: i dont want to handle errors at this level
+def upload_file(
+    client: AuthenticatedClient,
+    *,
+    mission_id: UUID,
+    filename: str,
+    path: Path,
+    verbose: bool = False,
+) -> UploadState:
+    """\
+    returns bytes uploaded
+    """
+
+    total_size = path.stat().st_size
+    with tqdm(
+        total=total_size,
+        unit="B",
+        unit_scale=True,
+        desc=f"uploading {path}...",
+        leave=False,
+        disable=not verbose,
+    ) as pbar:
+        endpoint = _get_s3_endpoint()
+
+        # get per file upload credentials
+        creds = _get_upload_creditials(
+            client, internal_filename=filename, mission_id=mission_id
+        )
+        if creds is None:
+            return UploadState.EXISTS
+
+        try:
+            _s3_upload(path, endpoint=endpoint, credentials=creds, pbar=pbar)
+        except Exception as e:
+            logger.error(format_traceback(e))
+            _cancel_file_upload(client, creds.file_id, mission_id)
+            return UploadState.CANCELED
+
+        else:
+            _confirm_file_upload(client, creds.file_id, b64_md5(path))
+            return UploadState.UPLOADED
+
+
+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 AccessDenied(
+            f"Failed to download file: {resp.json()['message']}"
+            f"Status Code: {resp.status_code}",
+        )
+
+    resp.raise_for_status()
+
+    return resp.text
+
+
+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}")
+
+    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,
+                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: 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)
+
+    return path.stat().st_size if state == UploadState.UPLOADED else 0
+
+
+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")

kleinkram-0.38.1.dev20241125112529/kleinkram/api/parsing.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import Any
+from typing import Dict
+from typing import Optional
+from uuid import UUID
+
+from kleinkram.errors import ParsingError
+from kleinkram.models import File
+from kleinkram.models import FileState
+from kleinkram.models import Mission
+from kleinkram.models import Project
+
+__all__ = [
+    "_parse_project",
+    "_parse_mission",
+    "_parse_file",
+]
+
+
+def _parse_project(project: Dict[str, Any]) -> Project:
+    try:
+        project_id = UUID(project["uuid"], version=4)
+        project_name = project["name"]
+        project_description = project["description"]
+
+        parsed = Project(
+            id=project_id, name=project_name, description=project_description
+        )
+    except Exception:
+        raise ParsingError(f"error parsing project: {project}")
+    return parsed
+
+
+def _parse_mission(
+    mission: Dict[str, Any], project: Optional[Project] = None
+) -> Mission:
+    try:
+        mission_id = UUID(mission["uuid"], version=4)
+        mission_name = mission["name"]
+
+        project_id = (
+            project.id if project else UUID(mission["project"]["uuid"], version=4)
+        )
+        project_name = project.name if project else mission["project"]["name"]
+
+        parsed = Mission(
+            id=mission_id,
+            name=mission_name,
+            project_id=project_id,
+            project_name=project_name,
+        )
+    except Exception:
+        raise ParsingError(f"error parsing mission: {mission}")
+    return parsed
+
+
+def _parse_file(file: Dict[str, Any], mission: Optional[Mission] = None) -> File:
+    try:
+        filename = file["filename"]
+        file_id = UUID(file["uuid"], version=4)
+        file_size = file["size"]
+        file_hash = file["hash"]
+
+        project_id = (
+            mission.project_id if mission else UUID(file["project"]["uuid"], version=4)
+        )
+        project_name = mission.project_name if mission else file["project"]["name"]
+
+        mission_id = mission.id if mission else UUID(file["mission"]["uuid"], version=4)
+        mission_name = mission.name if mission else file["mission"]["name"]
+
+        parsed = File(
+            id=file_id,
+            name=filename,
+            size=file_size,
+            hash=file_hash,
+            project_id=project_id,
+            project_name=project_name,
+            mission_id=mission_id,
+            mission_name=mission_name,
+            state=FileState(file["state"]),
+        )
+    except Exception:
+        raise ParsingError(f"error parsing file: {file}")
+    return parsed