antioch-py 2.0.6__py3-none-any.whl → 3.0.12__py3-none-any.whl

common/core/registry.py

@@ -1,12 +1,15 @@
 import json
+import os
+import tempfile
 from collections import defaultdict
 from datetime import datetime
 from pathlib import Path
 
-from common.ark import Ark as ArkDefinition, ArkReference, ArkVersionReference, AssetReference, AssetVersionReference
+from common.ark import Ark as ArkDefinition
 from common.constants import ANTIOCH_API_URL, get_ark_dir, get_asset_dir
-from common.core.auth import AuthError, AuthHandler
-from common.rome import RomeClient
+from common.core.auth import AuthHandler
+from common.core.rome import RomeClient
+from common.core.types import ArkReference, ArkVersionReference, AssetReference, AssetVersionReference
 
 
 def list_local_arks() -> list[ArkReference]:
@@ -17,17 +20,43 @@ def list_local_arks() -> list[ArkReference]:
     """
 
     arks_dir = get_ark_dir()
-    files_by_name = defaultdict(list)
-    for file_path in arks_dir.iterdir():
-        if file_path.is_file() and (file_path.name.endswith(":ark.json") or file_path.name.endswith(":asset.usdz")):
-            name = file_path.name.split(":")[0]
-            files_by_name[name].append(file_path)
 
+    # Group files by ark name and version
+    # File format: {name}:{version}:ark.json or {name}:{version}:asset.usdz
+    files_by_name_version: dict[str, dict[str, dict[str, Path]]] = defaultdict(lambda: defaultdict(dict))
+    for file_path in arks_dir.iterdir():
+        if not file_path.is_file():
+            continue
+        if file_path.name.endswith(":ark.json"):
+            name, version, _ = file_path.name.rsplit(":", 2)
+            files_by_name_version[name][version]["ark"] = file_path
+        elif file_path.name.endswith(":asset.usdz"):
+            name, version, _ = file_path.name.rsplit(":", 2)
+            files_by_name_version[name][version]["asset"] = file_path
+
+    # Build references for each ark
     results = []
-    for name, files in files_by_name.items():
-        ref = _build_ark_reference(name, files)
-        if ref is not None:
-            results.append(ref)
+    for name, versions in files_by_name_version.items():
+        version_refs = []
+        for version, files in versions.items():
+            ark_file = files.get("ark")
+            if ark_file is None:
+                continue
+            asset_file = files.get("asset")
+            ark_stat = ark_file.stat()
+            version_refs.append(
+                ArkVersionReference(
+                    version=version,
+                    full_path=str(ark_file),
+                    asset_path=str(asset_file) if asset_file else None,
+                    size_bytes=ark_stat.st_size,
+                    created_at=datetime.fromtimestamp(ark_stat.st_ctime).isoformat(),
+                    updated_at=datetime.fromtimestamp(ark_stat.st_mtime).isoformat(),
+                    asset_size_bytes=asset_file.stat().st_size if asset_file else None,
+                )
+            )
+        if version_refs:
+            results.append(build_ark_reference_from_versions(name, version_refs))
 
     return results
 
@@ -60,8 +89,8 @@ def get_ark_version_reference(name: str, version: str) -> ArkVersionReference:
             for version_ref in ark_ref.versions:
                 if version_ref.version == version:
                     return version_ref
-            raise FileNotFoundError(f"Version {version} of Ark {name} not found in local storage")
-    raise FileNotFoundError(f"No versions of Ark {name} found in local storage")
+            raise FileNotFoundError(f"Version {version} of Ark {name} not found locally. Please pull the Ark first.")
+    raise FileNotFoundError(f"No versions of Ark {name} found locally. Please pull the Ark first.")
 
 
 def get_asset_path(name: str, version: str, extension: str = "usdz", assert_exists: bool = True) -> Path:
@@ -94,19 +123,34 @@ def list_local_assets() -> list[AssetReference]:
     if not assets_dir.exists():
         return []
 
-    files_by_name: dict[str, list[Path]] = defaultdict(list)
+    # Group files by asset name and version
+    # File format: {name}:{version}:file.{extension}
+    files_by_name_version: dict[str, dict[str, Path]] = defaultdict(dict)
     for file_path in assets_dir.iterdir():
-        if file_path.is_file():
-            # Parse filename format: {name}:{version}:file.{extension}
-            parts = file_path.stem.split(":")
-            if len(parts) == 3 and parts[-1] == "file":
-                files_by_name[parts[0]].append(file_path)
+        if not file_path.is_file():
+            continue
+        parts = file_path.stem.split(":")
+        if len(parts) == 3 and parts[-1] == "file":
+            name, version = parts[0], parts[1]
+            files_by_name_version[name][version] = file_path
 
+    # Build references for each asset
     results = []
-    for name, files in files_by_name.items():
-        ref = _build_asset_reference(name, files)
-        if ref is not None:
-            results.append(ref)
+    for name, versions in files_by_name_version.items():
+        version_refs = []
+        for version, asset_file in versions.items():
+            asset_stat = asset_file.stat()
+            version_refs.append(
+                AssetVersionReference(
+                    version=version,
+                    full_path=str(asset_file),
+                    size_bytes=asset_stat.st_size,
+                    created_at=datetime.fromtimestamp(asset_stat.st_ctime).isoformat(),
+                    updated_at=datetime.fromtimestamp(asset_stat.st_mtime).isoformat(),
+                )
+            )
+        if version_refs:
+            results.append(build_asset_reference_from_versions(name, version_refs))
 
     return results
 
@@ -121,13 +165,8 @@ def list_remote_arks() -> list[ArkReference]:
     :raises AuthError: If not authenticated.
     """
 
-    # Get auth token
     auth = AuthHandler()
     token = auth.get_token()
-    if token is None:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and list arks
     rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
     return rome_client.list_arks()
 
@@ -136,6 +175,8 @@ def pull_remote_ark(name: str, version: str, overwrite: bool = False) -> ArkDefi
     """
     Pull an Ark from remote registry to local storage.
 
+    Downloads the Ark config (ark.json) and asset (asset.usdz) if present.
+
     Requires authentication.
 
     :param name: Name of the Ark.
@@ -148,30 +189,26 @@ def pull_remote_ark(name: str, version: str, overwrite: bool = False) -> ArkDefi
     # Check if Ark already exists locally
     arks_dir = get_ark_dir()
     ark_json_path = arks_dir / f"{name}:{version}:ark.json"
+    ark_asset_path = arks_dir / f"{name}:{version}:asset.usdz"
     if ark_json_path.exists() and not overwrite:
         return load_local_ark(name, version)
 
-    # Get auth token
     auth = AuthHandler()
     token = auth.get_token()
-    if not token:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and fetch Ark definition and save to local storage
     rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
-    ark = rome_client.get_ark(name=name, version=version)
-
-    # Save Ark JSON
-    with open(ark_json_path, "wb") as f:
-        f.write(json.dumps(ark).encode("utf-8"))
 
-    # Download asset usdz only if ark has asset_hash (hardware modules exist)
-    if ark.get("metadata", {}).get("asset_hash") is not None:
-        asset_content = rome_client.download_ark_assets(name=name, version=version)
-        with open(arks_dir / f"{name}:{version}:asset.usdz", "wb") as f:
-            f.write(asset_content)
-
-    return ArkDefinition(**ark)
+    print(f"Pulling {name} v{version}")
+    downloaded_asset = rome_client.pull_ark(
+        name=name,
+        version=version,
+        config_output_path=str(ark_json_path),
+        asset_output_path=str(ark_asset_path),
+    )
+    print("  ✓ Config downloaded")
+    if downloaded_asset:
+        print("  ✓ Asset downloaded")
+    print(f"✓ Ark {name} v{version} pulled successfully")
+    return load_local_ark(name, version)
 
 
 def list_remote_assets() -> list[AssetReference]:
@@ -184,17 +221,12 @@ def list_remote_assets() -> list[AssetReference]:
     :raises AuthError: If not authenticated.
     """
 
-    # Get auth token
     token = AuthHandler().get_token()
-    if token is None:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and list assets
     rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
     return rome_client.list_assets()
 
 
-def pull_remote_asset(name: str, version: str, overwrite: bool = False, show_progress: bool = True) -> Path:
+def pull_remote_asset(name: str, version: str, overwrite: bool = False) -> Path:
     """
     Pull an asset from remote registry to local storage.
 
@@ -203,129 +235,76 @@ def pull_remote_asset(name: str, version: str, overwrite: bool = False, show_pro
     :param name: Name of the asset.
     :param version: Version of the asset.
     :param overwrite: Overwrite local asset if it already exists.
-    :param show_progress: Show download progress bar.
     :return: Path to the downloaded asset file.
     :raises AuthError: If not authenticated.
     """
 
-    # Get auth token
-    token = AuthHandler().get_token()
-    if token is None:
-        raise AuthError("User not authenticated. Please login first")
-
-    # Create Rome client and get asset metadata to determine extension
-    rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
-    metadata = rome_client.get_asset_metadata(name=name, version=version)
-    extension = metadata.get("extension", "usdz")
-
     # Check if asset already exists locally
-    asset_file_path = get_asset_path(name=name, version=version, extension=extension, assert_exists=False)
+    # NOTE: Only checks USDZ assets for now
+    asset_file_path = get_asset_path(name=name, version=version, extension="usdz", assert_exists=False)
     if asset_file_path.exists() and not overwrite:
-        print(f"Asset {name}:{version} already exists locally, skipping download")
         return asset_file_path
 
-    # Download the asset file
-    rome_client.download_asset(name=name, version=version, output_path=str(asset_file_path), show_progress=show_progress)
-    return asset_file_path
+    token = AuthHandler().get_token()
+    rome_client = RomeClient(api_url=ANTIOCH_API_URL, token=token)
+    temp_path: str | None = None
+
+    try:
+        # Download to a temp file in the destination directory so publishing can be atomic
+        # This avoids EXDEV when the asset directory is a separate mount (common in Kubernetes)
+        asset_file_path.parent.mkdir(parents=True, exist_ok=True)
+        safe_prefix = f".{name.replace(':', '_')}.{version.replace(':', '_')}."
+        fd, temp_path = tempfile.mkstemp(prefix=safe_prefix, dir=str(asset_file_path.parent))
+        os.close(fd)
+
+        # Pull asset - metadata comes back from response body
+        metadata = rome_client.pull_asset(name=name, version=version, output_path=temp_path)
+        extension = metadata.get("extension", "usdz")
+
+        # Get final path with correct extension
+        asset_file_path = get_asset_path(name=name, version=version, extension=extension, assert_exists=False)
+        asset_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Publish atomically on the same filesystem
+        Path(temp_path).replace(asset_file_path)
+        return asset_file_path
+    finally:
+        # Clean up if download fails or publish raises
+        if temp_path is not None:
+            Path(temp_path).unlink(missing_ok=True)
 
 
-def _build_ark_reference(name: str, files: list[Path]) -> ArkReference | None:
+def build_ark_reference_from_versions(name: str, version_refs: list[ArkVersionReference]) -> ArkReference | None:
     """
-    Create an ArkReference from a list of files for a given ark.
+    Build an ArkReference from version references.
 
-    :param name: The name of the ark.
-    :param files: List of file paths (ark JSON and asset USDZ files).
-    :return: ArkReference object or None if no valid versions found.
+    :param name: The Ark name.
+    :param version_refs: List of ArkVersionReference instances.
+    :return: ArkReference or None if no versions exist.
     """
 
-    file_stats = [f.stat() for f in files]
-    created_at = min(datetime.fromtimestamp(stat.st_ctime) for stat in file_stats).isoformat()
-    updated_at = max(datetime.fromtimestamp(stat.st_mtime) for stat in file_stats).isoformat()
-
-    # Group files by version - parse from {name}:{version}:ark.json or {name}:{version}:asset.usdz
-    files_by_version: dict[str, list[Path]] = defaultdict(list)
-    for file_path in files:
-        files_by_version[file_path.name.split(":")[1]].append(file_path)
-
-    # Create an ArkVersionReference for each version
-    version_refs = []
-    for version, version_files in files_by_version.items():
-        ark_file = None
-        asset_file = None
-        for file_path in version_files:
-            if file_path.name.endswith(":ark.json"):
-                ark_file = file_path
-            elif file_path.name.endswith(":asset.usdz"):
-                asset_file = file_path
-        if ark_file is None:
-            continue
-
-        ark_stat = ark_file.stat()
-        version_refs.append(
-            ArkVersionReference(
-                version=version,
-                full_path=str(ark_file),
-                asset_path=str(asset_file) if asset_file else None,
-                size_bytes=ark_stat.st_size,
-                created_at=datetime.fromtimestamp(ark_stat.st_ctime).isoformat(),
-                updated_at=datetime.fromtimestamp(ark_stat.st_mtime).isoformat(),
-                asset_size_bytes=asset_file.stat().st_size if asset_file else None,
-            )
-        )
-
     if not version_refs:
         return None
 
-    return ArkReference(
-        name=name,
-        versions=version_refs,
-        created_at=created_at,
-        updated_at=updated_at,
-    )
+    # Aggregate timestamps from versions, fallback to empty string if all are missing
+    created_at = min((v.created_at for v in version_refs if v.created_at), default="")
+    updated_at = max((v.updated_at for v in version_refs if v.updated_at), default="")
+    return ArkReference(name=name, versions=version_refs, created_at=created_at, updated_at=updated_at)
 
 
-def _build_asset_reference(name: str, files: list[Path]) -> AssetReference | None:
+def build_asset_reference_from_versions(name: str, version_refs: list[AssetVersionReference]) -> AssetReference | None:
     """
-    Create an AssetReference from a list of files for a given asset.
+    Build an AssetReference from version references.
 
-    :param name: The name of the asset.
-    :param files: List of file paths (.usdz files).
-    :return: AssetReference object or None if no valid versions found.
+    :param name: The Asset name.
+    :param version_refs: List of AssetVersionReference instances.
+    :return: AssetReference or None if no versions exist.
     """
 
-    file_stats = [f.stat() for f in files]
-    created_at = min(datetime.fromtimestamp(stat.st_ctime) for stat in file_stats).isoformat()
-    updated_at = max(datetime.fromtimestamp(stat.st_mtime) for stat in file_stats).isoformat()
-
-    # Group files by version
-    files_by_version: dict[str, list[Path]] = defaultdict(list)
-    for file_path in files:
-        # Parse filename format: {name}:{version}:file.usdz
-        version = file_path.name.split(":")[1]
-        files_by_version[version].append(file_path)
-
-    # Create an AssetVersionReference for each version
-    version_refs = []
-    for version, version_files in files_by_version.items():
-        # Should only be one file per version, but take first if multiple
-        asset_file = version_files[0]
-        asset_stat = asset_file.stat()
-        version_refs.append(
-            AssetVersionReference(
-                version=version,
-                full_path=str(asset_file),
-                size_bytes=asset_stat.st_size,
-                created_at=datetime.fromtimestamp(asset_stat.st_ctime).isoformat(),
-                updated_at=datetime.fromtimestamp(asset_stat.st_mtime).isoformat(),
-            )
-        )
-
     if not version_refs:
         return None
 
-    return AssetReference(
-        name=name,
-        versions=version_refs,
-        created_at=created_at,
-        updated_at=updated_at,
-    )
+    # Aggregate timestamps from versions, fallback to empty string if all are missing
+    created_at = min((v.created_at for v in version_refs if v.created_at), default="")
+    updated_at = max((v.updated_at for v in version_refs if v.updated_at), default="")
+    return AssetReference(name=name, versions=version_refs, created_at=created_at, updated_at=updated_at)

common/core/rome.py

@@ -0,0 +1,251 @@
+from typing import overload
+
+import requests
+from requests import Response
+
+from common.core.types import ArkReference, AssetReference, TaskRun
+
+
+class RomeError(Exception):
+    """
+    Base error for Rome API operations.
+    """
+
+
+class RomeAuthError(RomeError):
+    """
+    Authentication error when interacting with Rome API.
+    """
+
+
+class RomeNetworkError(RomeError):
+    """
+    Network error when interacting with Rome API.
+    """
+
+
+class RomeClient:
+    """
+    Client for interacting with Rome (Antioch's cloud API).
+
+    Handles task runs, artifact uploads/downloads, and registry operations for Arks and Assets.
+    """
+
+    def __init__(self, api_url: str, token: str):
+        """
+        Initialize the Rome client.
+
+        :param api_url: Base URL for Rome API.
+        :param token: Authentication token.
+        """
+
+        self._api_url = api_url
+        self._token = token
+
+    def create_task_run(
+        self,
+        task_run: TaskRun,
+        upload_mcap: bool = False,
+        upload_bundle: bool = False,
+    ) -> tuple[str | None, str | None]:
+        """
+        Create a task run and optionally get signed URLs for artifact uploads.
+
+        :param task_run: TaskRun model with run data.
+        :param upload_mcap: Whether client will upload an MCAP file.
+        :param upload_bundle: Whether client will upload a bundle file.
+        :return: Tuple of (mcap_upload_url, bundle_upload_url).
+        """
+
+        response = self._send_request(
+            "POST",
+            "/tasks/runs",
+            json={
+                "task_run": task_run.model_dump(mode="json"),
+                "upload_mcap": upload_mcap,
+                "upload_bundle": upload_bundle,
+            },
+        )
+
+        return response.get("mcap_upload_url"), response.get("bundle_upload_url")
+
+    def list_arks(self) -> list[ArkReference]:
+        """
+        List all Arks from Rome registry.
+
+        :return: List of ArkReference objects from remote registry.
+        """
+
+        response = self._send_request("GET", "/ark/list")
+        return [ArkReference(**ark) for ark in response.get("data", [])]
+
+    def pull_ark(self, name: str, version: str, config_output_path: str, asset_output_path: str | None = None) -> bool:
+        """
+        Pull Ark config and optionally asset from Rome via signed URLs.
+
+        :param name: Name of the Ark.
+        :param version: Version of the Ark.
+        :param config_output_path: Path where ark.json should be saved.
+        :param asset_output_path: Path where asset.usdz should be saved (if present).
+        :return: True if an asset was downloaded, False otherwise.
+        """
+
+        if not self._token:
+            raise RomeAuthError("User not authenticated")
+
+        try:
+            # Get ark data with download URLs
+            response = self._send_request("GET", "/ark/get", json={"name": name, "version": version})
+
+            # Download config
+            config_response = requests.get(response["config_download_url"], timeout=60)
+            config_response.raise_for_status()
+            with open(config_output_path, "wb") as f:
+                f.write(config_response.content)
+
+            # Download asset if present and output path provided
+            has_asset = response.get("asset_download_url") is not None
+            if has_asset and asset_output_path:
+                asset_response = requests.get(response["asset_download_url"], timeout=None)
+                asset_response.raise_for_status()
+                with open(asset_output_path, "wb") as f:
+                    f.write(asset_response.content)
+                return True
+
+            return False
+        except requests.exceptions.RequestException as e:
+            raise RomeNetworkError(f"Network error: {e}") from e
+
+    def list_assets(self) -> list[AssetReference]:
+        """
+        List all assets from Rome registry.
+
+        :return: List of AssetReference objects from remote registry.
+        """
+
+        response = self._send_request("GET", "/asset/list")
+        return [AssetReference(**asset) for asset in response.get("data", [])]
+
+    def pull_asset(self, name: str, version: str, output_path: str) -> dict[str, str]:
+        """
+        Pull asset file from Rome registry via signed URL.
+
+        :param name: Name of the asset.
+        :param version: Version of the asset.
+        :param output_path: Path where the file should be saved.
+        :return: Metadata dictionary containing extension, file_size, and modified_time.
+        """
+
+        if not self._token:
+            raise RomeAuthError("User not authenticated")
+
+        try:
+            response = self._send_request("GET", "/asset/pull", params={"name": name, "version": version})
+            print(f"Downloading {name}:{version}...")
+            download_response = requests.get(response["download_url"], timeout=None)
+            download_response.raise_for_status()
+            with open(output_path, "wb") as f:
+                f.write(download_response.content)
+            return {
+                "extension": response.get("extension", ""),
+                "file_size": response.get("file_size", ""),
+                "modified_time": response.get("modified_time", ""),
+            }
+        except requests.exceptions.RequestException as e:
+            raise RomeNetworkError(f"Network error: {e}") from e
+
+    def get_gar_token(self) -> dict:
+        """
+        Get a GAR (Google Artifact Registry) access token for Docker operations.
+
+        :return: Dictionary with registry_host, repository, access_token, and expires_at.
+        """
+
+        response = self._send_request("GET", "/token/gar")
+        return response["data"]
+
+    @overload
+    def _send_request(
+        self,
+        method: str,
+        endpoint: str,
+        json: dict | None = None,
+        params: dict | None = None,
+        return_content: bool = False,
+    ) -> dict: ...
+
+    @overload
+    def _send_request(
+        self,
+        method: str,
+        endpoint: str,
+        json: dict | None = None,
+        params: dict | None = None,
+        return_content: bool = True,
+    ) -> bytes: ...
+
+    def _send_request(
+        self,
+        method: str,
+        endpoint: str,
+        json: dict | None = None,
+        params: dict | None = None,
+        return_content: bool = False,
+    ) -> dict | bytes:
+        """
+        Send a request to Rome API with standardized error handling.
+
+        :param method: HTTP method (GET, POST, etc.).
+        :param endpoint: API endpoint path.
+        :param json: Optional JSON payload.
+        :param params: Optional query parameters.
+        :param return_content: If True, return raw bytes content instead of JSON.
+        :return: Response JSON data or raw content bytes.
+        """
+
+        if not self._token:
+            raise RomeAuthError("User not authenticated")
+
+        try:
+            url = f"{self._api_url}{endpoint}"
+            headers = {"Authorization": f"Bearer {self._token}", "Content-Type": "application/json"}
+            response = requests.request(method, url, json=json, params=params, headers=headers, timeout=30)
+            self._check_response_errors(response)
+            if return_content:
+                return response.content
+
+            try:
+                return response.json()
+            except requests.exceptions.JSONDecodeError as e:
+                raise RomeError(f"Invalid JSON response: {e}") from e
+        except requests.exceptions.RequestException as e:
+            raise RomeNetworkError(f"Network error: {e}") from e
+
+    def _check_response_errors(self, response: Response) -> None:
+        """
+        Check response for HTTP errors and raise appropriate exceptions.
+
+        :param response: HTTP response object.
+        """
+
+        if response.status_code >= 400:
+            error_message = self._extract_error_message(response)
+            if response.status_code < 500:
+                raise RomeError(error_message)
+            raise RomeNetworkError(f"Server error: {error_message}")
+
+    def _extract_error_message(self, response: Response) -> str:
+        """
+        Extract error message from response JSON or return generic message.
+
+        :param response: HTTP response object.
+        :return: Error message string from response or generic HTTP status message.
+        """
+
+        try:
+            data = response.json()
+            if isinstance(data, dict) and "message" in data:
+                return data["message"]
+        except Exception:
+            pass
+        return f"HTTP {response.status_code}"