acceldata-sdk-python 26.7.0__py3-none-any.whl

acceldata/services/asset_service.py

@@ -0,0 +1,651 @@
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any, Mapping, Optional
+
+from pydantic import ValidationError
+
+from acceldata.client.http_client import BaseHttpClient
+from acceldata.client.transient_retry import (
+    RetryConfig,
+    effective_transient_retry,
+    transient_retry_call,
+)
+from acceldata.exceptions import APIError
+from acceldata.models.api.catalog.asset import Asset
+from acceldata.models.api.catalog.asset_label_request import AssetLabelRequest
+from acceldata.models.api.catalog.asset_tag_info import AssetTagInfo
+from acceldata.models.api.catalog.asset_tag_request import AssetTagRequest
+from acceldata.models.api.catalog.asset_tags_response import AssetTagsResponse
+from acceldata.models.api.catalog.catalog_asset_meta_data import CatalogAssetMetaData
+from acceldata.models.api.catalog.label import Label
+from acceldata.models.api.catalog.mini_profile_request import MiniProfileRequest
+from acceldata.models.api.catalog.profile_request import ProfileRequest
+from acceldata.models.api.catalog.sample_data_api_response import SampleDataApiResponse
+from acceldata.models.sdk.catalog.asset_resource import AssetResource
+from acceldata.models.sdk.catalog.profiling import AssetProfilingParams
+from acceldata.models.sdk.catalog.profile_request_resource import ProfileRequestResource
+
+logger = logging.getLogger("assetservice")
+
+# When sync=True, poll until the profile is clearly finished. The catalog can use
+# many non-terminal `status` values (PENDING, SCHEDULED, …) and may mirror progress
+# on `analysisJob.executionStatus`; a fixed in-progress set misses those and returns too early.
+_PROFILE_REQUEST_TERMINAL_STATUSES = frozenset(
+    {
+        "ABORTED",
+        "CANCELED",
+        "CANCELLED",
+        "COMPLETE",
+        "COMPLETED",
+        "ERRORED",
+        "EXECUTION_RESULT",
+        "FAILED",
+        "FINISHED",
+        "SKIPPED",
+        "SUCCESS",
+        "SUCCEEDED",
+        "SUCCESSFUL",
+        "WARNING",
+    }
+)
+_DATAPLANE_EXECUTION_TERMINAL_STATUSES = frozenset(
+    {
+        "ABORTED",
+        "EXECUTION_RESULT",
+        "FAILED",
+        "SKIPPED",
+        "SUCCESSFUL",
+    }
+)
+
+
+def _norm_catalog_enum(s: str) -> str:
+    """Uppercase and collapse spaces to underscores (``IN PROGRESS`` → ``IN_PROGRESS``)."""
+    return (s or "").upper().replace(" ", "_")
+
+
+def _profile_request_reached_terminal_state(pr: ProfileRequest) -> bool:
+    st = _norm_catalog_enum(pr.status or "")
+    if st in _PROFILE_REQUEST_TERMINAL_STATUSES:
+        return True
+    job = pr.analysis_job
+    if job is not None and job.execution_status:
+        ex = _norm_catalog_enum(job.execution_status or "")
+        if ex in _DATAPLANE_EXECUTION_TERMINAL_STATUSES:
+            return True
+    return False
+
+def _asset_list_from_response(response: dict[str, Any]) -> list[Any]:
+    for key in ("data", "assets"):
+        raw = response.get(key)
+        if raw is None:
+            continue
+        if isinstance(raw, (list, tuple)):
+            return list(raw)
+    raise KeyError(
+        "assets response must include 'data' or 'assets' with a list"
+    )
+
+
+class AssetService:
+    def __init__(self, http_client: BaseHttpClient) -> None:
+        self._client = http_client
+
+    def _catalog_body_error(
+        self,
+        *,
+        operation: str,
+        url: str,
+        response: object,
+        reason: str,
+    ) -> None:
+        """Raise :class:`APIError` when the catalog returns 2xx but the body is unusable."""
+        raise APIError(
+            message=f"{operation}: {reason}",
+            status_code=200,
+            response=response,
+            url=url,
+        )
+
+    @staticmethod
+    def _identifier_is_int_asset_id(identifier: str | int) -> bool:
+        """True when ``identifier`` should be resolved as a numeric asset id (not a uid string)."""
+        if isinstance(identifier, bool):
+            return False
+        if isinstance(identifier, int):
+            return True
+        if isinstance(identifier, str):
+            s = identifier.strip()
+            if s == "":
+                return False
+            try:
+                int(s)
+                return True
+            except ValueError:
+                return False
+        return False
+
+    def _asset_dicts_from_list_response(self, response: object) -> list[dict[str, Any]]:
+        """Read asset rows from ``GET /assets?...`` JSON (``data`` list or ``data.assets``)."""
+        if not isinstance(response, dict):
+            return []
+        data = response.get("data")
+        logger.debug(f"data: {data}")
+        if isinstance(data, list):
+            return [x for x in data if isinstance(x, dict)]
+        if isinstance(data, dict):
+            nested = data.get("assets")
+            if isinstance(nested, list):
+                return [x for x in nested if isinstance(x, dict)]
+        return []
+
+    def _wrap_asset(self, asset: Asset) -> AssetResource:
+        return AssetResource(asset_service=self, model=asset)
+
+    def _first_asset_dict(self, response: object) -> dict[str, Any] | None:
+        rows = self._asset_dicts_from_list_response(response)
+        return rows[0] if rows else None
+
+    def get_asset(self, identifier: str | int) -> AssetResource:
+        """
+        Find an asset by uid or numeric id.
+
+        :param identifier: Asset uid, or numeric id (``int`` or decimal string).
+        """
+        if identifier is None or (isinstance(identifier, str) and identifier.strip() == ""):
+            raise ValueError("asset identifier is required")
+        logger.debug("get_asset identifier=%r", identifier)
+        if self._identifier_is_int_asset_id(identifier):
+            aid = int(identifier) if isinstance(identifier, str) else int(identifier)
+            return self.get_asset_by_id(aid)
+        return self.get_asset_by_uid(str(identifier).strip())
+
+    def get_asset_by_id(self, id: int) -> AssetResource:
+        if id is None:
+            raise ValueError("Asset id is required")
+        logger.debug("get_asset_by_id id=%s", id)
+        url = self._client.catalog_urls.assets_by_ids(asset_ids=id)
+        response = self._client._get(url)
+        logger.debug(f"get_asset_by_id response: {response}")
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation="get_asset_by_id",
+                url=url,
+                response=response,
+                reason="expected JSON object response",
+            )
+        row = self._first_asset_dict(response)
+        asset = Asset.from_dict(row)
+        if asset is None:
+            self._catalog_body_error(
+                operation="get_asset_by_id",
+                url=url,
+                response=response,
+                reason="no asset in response body",
+            )
+        logger.debug("get_asset_by_id: parsed asset id=%s asset_id=%s", id, asset.id)
+        return self._wrap_asset(asset)
+
+    def get_asset_by_uid(self, uid: str) -> AssetResource:
+        url = self._client.catalog_urls.assets_by_uid(uid=uid)
+        response = self._client._get(url)
+        logger.debug(f"get_asset_by_uid :{response}")
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation="get_asset_by_uid",
+                url=url,
+                response=response,
+                reason="expected JSON object response",
+            )
+
+        row = self._first_asset_dict(response)
+        logger.debug(f"asset row: {row}")
+        asset = Asset.from_dict(row)
+        if asset is None:
+            self._catalog_body_error(
+                operation="get_asset_by_uid",
+                url=url,
+                response=response,
+                reason="no asset in response body",
+            )
+        return self._wrap_asset(asset)
+
+
+    def get_asset_metadata(self, asset_id: int) -> CatalogAssetMetaData:
+        logger.debug("get_asset_metadata asset_id=%s", asset_id)
+        url = self._client.catalog_urls.asset_metadata(asset_id=asset_id)
+        response = self._client._get(url=url)
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation="get_asset_metadata",
+                url=url,
+                response=response,
+                reason="expected JSON object response",
+            )
+        data = response.get("data")
+        logger.debug("asset meta data: %s", data)
+        if data is None:
+            self._catalog_body_error(
+                operation="get_asset_metadata",
+                url=url,
+                response=response,
+                reason="missing 'data' in response body",
+            )
+        meta = CatalogAssetMetaData.from_dict(data)
+        if meta is None:
+            self._catalog_body_error(
+                operation="get_asset_metadata",
+                url=url,
+                response=response,
+                reason="could not parse catalog asset metadata",
+            )
+        return meta
+
+
+    def add_custom_metadata(
+            self,
+            asset_id: int,
+            custom_metadata: Mapping[str, str],
+    ) -> CatalogAssetMetaData:
+        """POST custom metadata as string key/value pairs (JSON object with string values)."""
+        logger.debug(
+            "add_custom_metadata asset_id=%s keys=%s",
+            asset_id,
+            list(custom_metadata.keys()),
+        )
+        url = self._client.catalog_urls.asset_custom_metadata(asset_id=asset_id)
+        response = self._client._post(url, payload=dict(custom_metadata))
+        logger.debug(response)
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation="add_custom_metadata",
+                url=url,
+                response=response,
+                reason="expected JSON object response",
+            )
+        data = response.get("data")
+        logger.debug("asset meta data: %s", data)
+        if data is None:
+            self._catalog_body_error(
+                operation="add_custom_metadata",
+                url=url,
+                response=response,
+                reason="missing 'data' in response body",
+            )
+        meta = CatalogAssetMetaData.from_dict(data)
+        if meta is None:
+            self._catalog_body_error(
+                operation="add_custom_metadata",
+                url=url,
+                response=response,
+                reason="could not parse catalog asset metadata",
+            )
+        return meta
+
+    def _profile_request_from_catalog_response(
+            self,
+            response: object,
+            *,
+            op: str,
+            asset_id: int,
+            url: str,
+    ) -> ProfileRequest:
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation=op,
+                url=url,
+                response=response,
+                reason=f"non-dict response (asset_id={asset_id})",
+            )
+        # POST /assets/{id}/profile returns StandardResponse_ProfileRequest (``data``).
+        # GET /assets/{id}/profile/{reqId} returns ProfileRequestResponse (``profileRequest``).
+        data = response.get("data")
+        if data is None:
+            data = response.get("profileRequest")
+        result = ProfileRequest.from_dict(data)
+        if result is None:
+            self._catalog_body_error(
+                operation=op,
+                url=url,
+                response=response,
+                reason=f"missing or invalid profile request data (asset_id={asset_id})",
+            )
+        return result
+
+    def _mini_profile_request_from_catalog_response(
+        self,
+        response: object,
+        *,
+        op: str,
+        asset_id: int,
+        url: str,
+    ) -> MiniProfileRequest:
+        """Parse ``StandardResponse_MiniProfileRequest`` (GET latest profile on ``/assets/{id}/profile``)."""
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation=op,
+                url=url,
+                response=response,
+                reason=f"non-dict response (asset_id={asset_id})",
+            )
+        data = response.get("data")
+        result = MiniProfileRequest.from_dict(data)
+        if result is None:
+            self._catalog_body_error(
+                operation=op,
+                url=url,
+                response=response,
+                reason=f"missing or invalid mini profile request data (asset_id={asset_id})",
+            )
+        return result
+
+    def profile_asset(
+        self,
+        asset_id: int,
+        profiling_params: Optional[AssetProfilingParams] = None,
+        *,
+        sync: bool = False,
+        transient_retry: RetryConfig | None = None,
+    ) -> ProfileRequestResource:
+        """POST ``/assets/{id}/profile`` with a catalog profiling body (``profile_asset`` parity).
+
+        When ``transient_retry`` is set, the POST retries transient HTTP/transport failures
+        before surfacing an error. HTTP 409 is included in the default transient
+        ``retry_status_codes``.
+
+        When ``sync`` is ``True``, this method polls
+        ``GET /assets/{id}/profile/{reqId}`` until the top-level request ``status`` or, when
+        present, the nested ``analysisJob.executionStatus`` reflects a **terminal** state
+        (e.g. ``SUCCESSFUL`` / ``FAILED``), not only a fixed in-progress set.
+        """
+        if profiling_params is None:
+            raise TypeError("profiling_params is required for profile_asset")
+        logger.debug("Marker config: %s", profiling_params.marker_config)
+        logger.debug("Converting profiling_params to dictionary")
+        start_profiling_request = profiling_params.to_start_profiling_request()
+        payload = start_profiling_request.to_dict() or {}
+        logger.debug("StartProfilingRequest to dictionary: %s", payload)
+        logger.debug("Executing Profiling with the following payload: %s", payload)
+        url = self._client.catalog_urls.asset_profile(asset_id=asset_id)
+        tr_cfg = effective_transient_retry(transient_retry, None)
+        try:
+            response = transient_retry_call(
+                lambda: self._client._post(url, payload=payload or {}),
+                tr_cfg,
+                logger=logger,
+                operation_name="profile_asset",
+            )
+            logger.debug(f"profile_asset response: {response}")
+            model = self._profile_request_from_catalog_response(
+                response, op="profile_asset", asset_id=asset_id, url=url
+            )
+            if sync:
+                model = self._poll_profile_request_until_terminal(
+                    asset_id=asset_id, req_id=int(model.id)
+                )
+            status = model.status or "UNKNOWN"
+            if sync:
+                message = (
+                    f"Profiling completed for asset {asset_id} with request {model.id} "
+                    f"(status={status})."
+                )
+            else:
+                message = (
+                    f"Profiling triggered for asset {asset_id} with request {model.id} "
+                    f"(current status={status})."
+                )
+            return ProfileRequestResource(self, model, operation_message=message)
+        except APIError as exc:
+            logger.error(
+                "profile_asset: request failed for asset_id=%s (HTTP %s): %s",
+                asset_id,
+                exc.status_code,
+                exc.message,
+            )
+            raise
+
+    def _poll_profile_request_until_terminal(self, *, asset_id: int, req_id: int) -> ProfileRequest:
+        """Poll profile request details until top-level or dataplane status is terminal."""
+        profile_request = self.get_profile_request_details(asset_id=asset_id, req_id=req_id)
+        while not _profile_request_reached_terminal_state(profile_request):
+            time.sleep(5)
+            profile_request = self.get_profile_request_details(asset_id=asset_id, req_id=req_id)
+        return profile_request
+
+    def cancel_profile(self, profile_req_id: int):
+        return self._client.cancel_profile(profile_req_id)
+
+    def get_asset_tags(self, asset_id: int) -> list[AssetTagInfo]:
+        """GET tags for one asset (``.../assets/{id}/tags``)."""
+        if asset_id is None:
+            raise ValueError("Asset id is required")
+        logger.debug("get_asset_tags asset_id=%s", asset_id)
+        url = self._client.catalog_urls.asset_tags(asset_id=asset_id)
+        response = self._client._get(url=url)
+        return self._asset_tags_from_catalog_response(
+            response, operation="get_asset_tags", url=url
+        )
+
+    def add_asset_tag(self, asset_id: int, payload: AssetTagRequest) -> AssetTagInfo:
+        """POST a tag on an asset (``.../assets/{id}/tag``)."""
+        if asset_id is None:
+            raise ValueError("Asset id is required")
+        logger.debug("add_asset_tag asset_id=%s", asset_id)
+        url = self._client.catalog_urls.asset_tag(asset_id=asset_id)
+        response = self._client._post(url, payload=payload.to_dict())
+        if not isinstance(response, dict):
+            self._catalog_body_error(
+                operation="add_asset_tag",
+                url=url,
+                response=response,
+                reason="expected JSON object response",
+            )
+        tag = AssetTagInfo.from_dict(response)
+        if tag is None:
+            self._catalog_body_error(
+                operation="add_asset_tag",
+                url=url,
+                response=response,
+                reason="could not parse asset tag from response body",
+            )
+        return tag
+
+    def add_asset_labels(self, asset_id: int, payload: AssetLabelRequest) -> list[Label]:
+        """PUT labels on an asset (``.../assets/{id}/labels``)."""
+        if asset_id is None:
+            raise ValueError("Asset id is required")
+        logger.debug("add_asset_labels asset_id=%s", asset_id)
+        url = self._client.catalog_urls.asset_labels(asset_id=asset_id)
+        response = self._client._put(url, payload=payload.to_dict())
+        return self._labels_from_catalog_response(
+            response, operation="add_asset_labels", url=url
+        )
+
+    def get_asset_labels(self, asset_id: int) -> list[Label]:
+        """GET labels for one asset (``.../assets/{id}/labels``)."""
+        if asset_id is None:
+            raise ValueError("Asset id is required")
+        logger.debug("get_asset_labels asset_id=%s", asset_id)
+        url = self._client.catalog_urls.asset_labels(asset_id=asset_id)
+        response = self._client._get(url=url)
+        return self._labels_from_catalog_response(
+            response, operation="get_asset_labels", url=url
+        )
+
+    @staticmethod
+    def _asset_tags_from_catalog_response(
+        response: object, *, operation: str, url: str
+    ) -> list[AssetTagInfo]:
+        if isinstance(response, dict):
+            inner = response.get("data")
+            if isinstance(inner, dict):
+                response = inner
+            try:
+                parsed = AssetTagsResponse.from_dict(response)
+                if parsed is not None and parsed.tags is not None:
+                    return list(parsed.tags)
+            except (ValidationError, ValueError, TypeError):
+                pass
+            for key in ("tags", "assetTags"):
+                raw = response.get(key)
+                if isinstance(raw, list):
+                    tags: list[AssetTagInfo] = []
+                    for item in raw:
+                        if not isinstance(item, dict):
+                            continue
+                        tag = AssetTagInfo.from_dict(item)
+                        if tag is not None:
+                            tags.append(tag)
+                    return tags
+        raise APIError(
+            message=f"{operation}: missing or invalid tags in response body",
+            status_code=200,
+            response=response,
+            url=url,
+        )
+
+    @staticmethod
+    def _labels_from_catalog_response(
+        response: object, *, operation: str, url: str
+    ) -> list[Label]:
+        raw: list[Any] | None = None
+        if isinstance(response, list):
+            raw = list(response)
+        elif isinstance(response, dict):
+            data = response.get("data")
+            if isinstance(data, list):
+                raw = list(data)
+            elif isinstance(data, dict) and isinstance(data.get("labels"), list):
+                raw = list(data["labels"])
+            elif isinstance(response.get("labels"), list):
+                raw = list(response["labels"])
+        if raw is None:
+            raise APIError(
+                message=f"{operation}: expected a JSON array of labels",
+                status_code=200,
+                response=response,
+                url=url,
+            )
+        labels: list[Label] = []
+        for item in raw:
+            if not isinstance(item, dict):
+                continue
+            label = Label.from_dict(item)
+            if label is not None:
+                labels.append(label)
+        return labels
+
+    def get_tags(self) -> AssetTagsResponse | None:
+        logger.debug("get_tags (global)")
+        url = self._client.catalog_urls.assets_tags()
+        response = self._client._get(url=url)
+        if not isinstance(response, dict):
+            return None
+        inner = response.get("data")
+        if isinstance(inner, dict):
+            try:
+                return AssetTagsResponse.from_dict(inner)
+            except (ValidationError, ValueError, TypeError):
+                pass
+        try:
+            return AssetTagsResponse.from_dict(response)
+        except (ValidationError, ValueError, TypeError):
+            logger.debug("get_tags: AssetTagsResponse.from_dict failed", exc_info=True)
+            return None
+
+    def get_latest_profile_status(self, asset_id: int) -> MiniProfileRequest:
+        """GET latest profile status for the asset (same path as profile POST; body is ``MiniProfileRequest``)."""
+        logger.debug("get_latest_profile_status asset_id=%s", asset_id)
+        url = self._client.catalog_urls.asset_profile(asset_id=asset_id)
+        response = self._client._get(url=url)
+        return self._mini_profile_request_from_catalog_response(
+            response, op="get_latest_profile_status", asset_id=asset_id, url=url
+        )
+
+    def get_profile_request_details(self, asset_id: int, req_id: int) -> ProfileRequest:
+        """GET one profile request by asset id and profile request id."""
+        logger.debug(
+            "get_profile_request_details asset_id=%s req_id=%s", asset_id, req_id
+        )
+        url = self._client.catalog_urls.profile_request_details(
+            asset_id=asset_id, req_id=req_id
+        )
+        response = self._client._get(url=url)
+        return self._profile_request_from_catalog_response(
+            response,
+            op="get_profile_request_details",
+            asset_id=asset_id,
+            url=url,
+        )
+
+
+    def get_sample_data_result(
+        self,
+        request_id: str | SampleDataApiResponse | None,
+    ) -> SampleDataApiResponse | None:
+        if isinstance(request_id, SampleDataApiResponse):
+            initial = request_id
+            rid = initial.request_id
+            if not rid:
+                logger.debug(
+                    "get_sample_data_result: inline sample response job_execution_status=%s",
+                    initial.job_execution_status,
+                )
+                return initial
+            request_id = rid
+
+        if not request_id:
+            raise ValueError("request_id is null")
+
+        logger.debug("get_sample_data_result request_id=%s", request_id)
+        url = self._client.catalog_urls.asset_sample_result(request_id=request_id)
+        response = self._client._get(url=url)
+        if not isinstance(response, dict):
+            logger.error("get_sample_data_result: non-dict response request_id=%s", request_id)
+        return SampleDataApiResponse.from_dict(response)
+
+    def sample_data(
+        self,
+        asset_id: int,
+        *,
+        sync: bool = True,
+    ) -> SampleDataApiResponse | None:
+        """Trigger sample data and optionally poll for completion.
+
+        When ``sync`` is ``True`` (default), this method blocks and polls
+        ``get_sample_data_result`` until ``jobExecutionStatus`` is not
+        ``IN_PROGRESS``. When ``sync`` is ``False``, it returns the initial
+        async response (typically containing ``request_id``) without polling.
+        """
+        logger.debug("sample_data asset_id=%s sync=%s", asset_id, sync)
+        url = self._client.catalog_urls.asset_sample_async(asset_id=asset_id)
+        response = self._client._post(url, payload={})
+        if not isinstance(response, dict):
+            logger.debug("sample_data: non-dict async response asset_id=%s", asset_id)
+            return None
+        try:
+            initial = SampleDataApiResponse.from_dict(response)
+        except (ValidationError, ValueError, TypeError):
+            logger.debug("sample_data: parse async response failed asset_id=%s", asset_id, exc_info=True)
+            return None
+        if initial is None:
+            return None
+        rid = initial.request_id
+        if not rid:
+            logger.debug("sample_data: no request_id in async response asset_id=%s", asset_id)
+            return initial
+        if not sync:
+            return initial
+
+        sample_data_result = self.get_sample_data_result(rid)
+        while (
+            sample_data_result is not None
+            and (sample_data_result.job_execution_status or "") == "IN_PROGRESS"
+        ):
+            time.sleep(5)
+            sample_data_result = self.get_sample_data_result(rid)
+        return sample_data_result
+
+