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

acceldata/services/pipeline_service.py

@@ -0,0 +1,494 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import asdict, is_dataclass
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Optional
+
+from acceldata.models.api.pipeline.pipeline_runs_response import PipelineRunsResponse
+from pydantic import BaseModel, JsonValue
+
+from acceldata.client.http_client import BaseHttpClient
+from acceldata.exceptions import InvalidPipelineRunLookupError, PipelineRunNotFoundError
+from acceldata.models.api.pipeline.node_response import NodeResponse
+from acceldata.models.api.pipeline.pipeline_details_response import PipelineDetailsResponse
+from acceldata.models.api.pipeline.pipeline_response import PipelineResponse
+from acceldata.models.api.pipeline.pipeline import Pipeline
+from acceldata.models.api.pipeline.pipeline_run import PipelineRun
+from acceldata.models.api.pipeline.tag import Tag
+from acceldata.models.api.pipeline.pipelines_listing_response import PipelinesListingResponse
+from acceldata.models.api.pipeline.span import Span
+from acceldata.models.api.pipeline.span_event import SpanEvent
+from acceldata.models.api.pipeline.span_event_request import SpanEventRequest
+from acceldata.models.api.pipeline.span_request import SpanRequest
+from acceldata.models.api.pipeline.spans_response import SpansResponse
+from acceldata.models.sdk.pipeline.create_pipeline_input_request import CreatePipelineInputRequest
+from acceldata.models.sdk.pipeline.pipeline_run_input import PipelineRunResult, PipelineRunStatus
+from acceldata.services.pipeline_response_parse import (
+    pipeline_response_from_api_dict,
+    span_from_api_dict,
+    span_request_from_api_dict,
+    spans_response_from_api_dict,
+)
+from acceldata.utils.operation_ack import ack_failure, ack_success
+
+if TYPE_CHECKING:
+    from acceldata.models.sdk.pipeline.span_resource import SpanResource
+
+logger = logging.getLogger("pipelineservice")
+
+
+def _sanitize_for_json(obj: Any) -> Any:
+    if obj is None:
+        return None
+    if isinstance(obj, Enum):
+        return obj.value
+    if isinstance(obj, datetime):
+        return obj.isoformat()
+    if isinstance(obj, dict):
+        return {str(k): _sanitize_for_json(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_sanitize_for_json(x) for x in obj]
+    if isinstance(obj, (str, int, float, bool)):
+        return obj
+    if isinstance(obj, BaseModel):
+        if hasattr(obj, "model_dump"):
+            dumped = obj.model_dump(by_alias=True, exclude_none=True)
+        else:
+            dumped = obj.dict(by_alias=True, exclude_none=True)
+        return _sanitize_for_json(dumped)
+    if is_dataclass(obj):
+        return _sanitize_for_json(asdict(obj))
+    return obj
+
+
+class PipelineService:
+    def __init__(self, http_client: BaseHttpClient) -> None:
+        self._client = http_client
+        self.logger = logger
+
+    def _convert_pipeline_to_dict(self, pipeline: Optional[CreatePipelineInputRequest]) -> dict[str, JsonValue]:
+        if isinstance(pipeline, dict):
+            return dict(pipeline)
+        if pipeline is None:
+            raise TypeError("create_pipeline_input_request is required")
+        # CreatePipelineInputRequest is a Pydantic BaseModel, not a dataclass — use to_dict().
+        if isinstance(pipeline, BaseModel):
+            outer = pipeline.to_dict()
+            inner = outer.get("pipeline")
+            if not isinstance(inner, dict):
+                raise TypeError(
+                    "CreatePipelineInputRequest must serialize to a dict with a 'pipeline' object"
+                )
+            raw = _sanitize_for_json(inner)
+        elif is_dataclass(pipeline):
+            raw = _sanitize_for_json(asdict(pipeline))
+        else:
+            raise TypeError(
+                f"Unsupported pipeline request type: {type(pipeline).__name__}"
+            )
+        if not isinstance(raw, dict):
+            raise TypeError("Pipeline serialization must produce a JSON object")
+        return raw
+
+    def create_pipeline(self, create_pipeline_input_request: Optional[CreatePipelineInputRequest] = None,
+                        **kwargs) -> PipelineResponse:
+        payload = self._convert_pipeline_to_dict(create_pipeline_input_request)
+        url = self._client.pipeline_urls.pipelines_root()
+        response = self._client._put(url, payload={"pipeline": payload})
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from create pipeline, got {type(response).__name__}")
+        p = response.get("pipeline")
+        if not isinstance(p, dict):
+            raise TypeError("Create pipeline response must include a 'pipeline' object")
+        out = pipeline_response_from_api_dict(response)
+        pid = out.pipeline.id if out.pipeline is not None else None
+        ack_success(self.logger, "create_pipeline", pipeline_id=pid)
+        return out
+
+    def update_pipeline(self, pipeline: Pipeline) -> PipelineResponse:
+        """PUT ``/pipelines`` with a full :class:`~acceldata.models.api.pipeline.pipeline.Pipeline` body (must include ``id`` when updating)."""
+        if pipeline.id is None:
+            raise ValueError("Pipeline id is required to update a pipeline")
+        raw = _sanitize_for_json(pipeline.to_dict())
+        if not isinstance(raw, dict):
+            raise TypeError("Pipeline serialization must produce a JSON object")
+        url = self._client.pipeline_urls.pipelines_root()
+        response = self._client._put(url, payload={"pipeline": raw})
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from update pipeline, got {type(response).__name__}")
+        p = response.get("pipeline")
+        if not isinstance(p, dict):
+            raise TypeError("Update pipeline response must include a 'pipeline' object")
+        out = pipeline_response_from_api_dict(response)
+        pid = out.pipeline.id if out.pipeline is not None else None
+        ack_success(self.logger, "update_pipeline", pipeline_id=pid)
+        return out
+
+    def replace_pipeline_tags(
+        self, pipeline_identity: str | int, tags: list[Tag]
+    ) -> PipelineResponse:
+        """GET the pipeline by identity, replace its ``tags`` list, and PUT the full pipeline back."""
+        current = self.get_pipeline(str(pipeline_identity))
+        pl = current.pipeline
+        if pl is None:
+            raise TypeError("get_pipeline returned no pipeline object")
+        updated = pl.model_copy(update={"tags": list(tags)})
+        return self.update_pipeline(updated)
+
+    def delete_pipeline(self, pipeline_id: str | int) -> dict[str, JsonValue]:
+        url = self._client.pipeline_urls.pipeline_by_identity(pipeline_id)
+        response = self._client._delete(url)
+        ack_success(self.logger, "delete_pipeline", pipeline_id=pipeline_id)
+        if isinstance(response, dict):
+            return response
+        return {}
+
+    def create_job(self, job: dict[str, Any], pipeline_id: int) -> NodeResponse:
+        url = self._client.pipeline_urls.pipeline_jobs(pipeline_id)
+        response = self._client._put(url, payload=job)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from create job, got {type(response).__name__}")
+        node = NodeResponse.from_dict(response)
+        job_uid = node.node.uid if node.node is not None else None
+        ack_success(self.logger, "create_job", pipeline_id=pipeline_id, job_uid=job_uid)
+        return node
+
+    def get_pipeline_run(
+            self,
+            pipeline_run_id: Optional[int] = None,
+            continuation_id: Optional[str] = None,
+            pipeline_id: Optional[int] = None,
+    ) -> PipelineRun:
+        if pipeline_run_id is not None:
+            url = self._client.pipeline_urls.pipeline_run_by_id(pipeline_run_id)
+            response = self._client._get(url)
+            if not isinstance(response, dict):
+                raise TypeError(f"Expected dict from get pipeline run, got {type(response).__name__}")
+            r = response.get("run")
+        elif continuation_id is not None and pipeline_id is not None:
+            url = self._client.pipeline_urls.pipeline_runs(pipeline_id)
+            response = self._client._get(url, params={"continuationId": continuation_id})
+            if not isinstance(response, dict):
+                raise TypeError(f"Expected dict from pipeline runs listing, got {type(response).__name__}")
+            runs = response.get("runs") or []
+            if not isinstance(runs, list) or not runs:
+                ack_failure(
+                    self.logger,
+                    "get_pipeline_run",
+                    "run_not_found",
+                    continuation_id=continuation_id,
+                    pipeline_id=pipeline_id,
+                )
+                raise PipelineRunNotFoundError(
+                    (
+                        f"No pipeline run found for continuation_id={continuation_id!r} on pipeline_id={pipeline_id}. "
+                        "Confirm the continuation id matches the active run, or use get_latest_pipeline_run() "
+                        "when a single latest run is sufficient."
+                    ),
+                    continuation_id=continuation_id,
+                    pipeline_id=pipeline_id,
+                )
+            r = runs[0]
+        else:
+            ack_failure(
+                self.logger,
+                "get_pipeline_run",
+                "invalid_lookup",
+                pipeline_run_id=pipeline_run_id,
+                continuation_id=continuation_id,
+                pipeline_id=pipeline_id,
+            )
+            raise InvalidPipelineRunLookupError(
+                "Provide pipeline_run_id, or both continuation_id and pipeline_id together."
+            )
+        if not isinstance(r, dict):
+            raise TypeError("Pipeline run response must include a 'run' object")
+        run = PipelineRun.from_dict(r)
+        ack_success(
+            self.logger,
+            "get_pipeline_run",
+            pipeline_run_id=run.id,
+            continuation_id=continuation_id,
+        )
+        return run
+
+    def get_latest_pipeline_run(self, pipeline_id: int) -> PipelineRun:
+        url = self._client.pipeline_urls.pipeline_latest_run(pipeline_id)
+        response = self._client._get(url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from latest pipeline run, got {type(response).__name__}")
+        r = response.get("run")
+        if not isinstance(r, dict):
+            raise TypeError("Latest pipeline run response must include a 'run' object")
+        run = PipelineRun.from_dict(r)
+        ack_success(self.logger, "get_latest_pipeline_run", pipeline_id=pipeline_id, pipeline_run_id=run.id)
+        return run
+
+    def get_pipeline_runs(self, pipeline_id: int) -> PipelineRunsResponse:
+        url = self._client.pipeline_urls.pipeline_runs(pipeline_id)
+        response = self._client._get(url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from pipeline runs, got {type(response).__name__}")
+        listing = PipelineRunsResponse.from_dict(response)
+        ack_success(self.logger, "get_pipeline_runs", pipeline_id=pipeline_id)
+        return listing
+
+    def create_pipeline_run(
+            self,
+            pipeline_id: int,
+            *,
+            context_data: Optional[dict[str, Any]] = None,
+            continuation_id: Optional[str] = None,
+            started_at: Optional[datetime] = None,
+    ) -> PipelineRun:
+        run_inner: dict[str, Any] = {
+            "pipelineId": pipeline_id,
+            "status": PipelineRunStatus.STARTED.value,
+        }
+        if continuation_id is not None:
+            run_inner["continuationId"] = continuation_id
+        if context_data is not None:
+            run_inner["args"] = context_data
+        if started_at is not None:
+            run_inner["startedAt"] = started_at.isoformat()
+        payload = {"run": _sanitize_for_json(run_inner)}
+        url = self._client.pipeline_urls.pipeline_runs(pipeline_id)
+        response = self._client._post(url, payload=payload)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from create pipeline run, got {type(response).__name__}")
+        r = response.get("run")
+        if not isinstance(r, dict):
+            raise TypeError("Create pipeline run response must include a 'run' object")
+        run = PipelineRun.from_dict(r)
+        ack_success(
+            self.logger,
+            "create_pipeline_run",
+            pipeline_id=pipeline_id,
+            pipeline_run_id=run.id,
+            continuation_id=continuation_id,
+        )
+        return run
+
+    def update_pipeline_run(
+            self,
+            pipeline_run_id: int,
+            *,
+            pipeline_id: int,
+            result: PipelineRunResult,
+            status: PipelineRunStatus,
+            version_id: Optional[int] = None,
+            context_data: Optional[dict[str, Any]] = None,
+            continuation_id: Optional[str] = None,
+            started_at: Optional[datetime] = None,
+            finished_at: Optional[datetime] = None,
+    ) -> PipelineRun:
+
+        run_inner: dict[str, Any] = {
+            "pipelineId": pipeline_id,
+            "result": result.value,
+            "status": status.value,
+        }
+        if version_id is not None:
+            run_inner["versionId"] = version_id
+        if continuation_id is not None:
+            run_inner["continuationId"] = continuation_id
+        if context_data is not None:
+            run_inner["args"] = context_data
+        if started_at is not None:
+            run_inner["startedAt"] = started_at.isoformat()
+        if finished_at is not None:
+            run_inner["finishedAt"] = finished_at.isoformat()
+        payload = {"run": _sanitize_for_json(run_inner)}
+        url = self._client.pipeline_urls.pipeline_run_by_id(pipeline_run_id)
+        response = self._client._put(url, payload=payload)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from update pipeline run, got {type(response).__name__}")
+        r = response.get("run")
+        if not isinstance(r, dict):
+            raise TypeError("Update pipeline run response must include a 'run' object")
+        run = PipelineRun.from_dict(r)
+        ack_success(
+            self.logger,
+            "update_pipeline_run",
+            pipeline_run_id=pipeline_run_id,
+            pipeline_id=pipeline_id,
+            result=result.value,
+            status=status.value,
+        )
+        return run
+
+    def create_span(
+            self,
+            pipeline_run_id: int,
+            payload: dict[str, JsonValue],
+    ) -> Span:
+        url = self._client.pipeline_urls.pipeline_run_spans(pipeline_run_id)
+        response = self._client._post(url, payload=payload)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from create span, got {type(response).__name__}")
+        s = response.get("span")
+        if not isinstance(s, dict):
+            raise TypeError("Create span response must include a 'span' object")
+        res = span_from_api_dict(s)
+        if res is None:
+            raise TypeError("Could not parse span from create span response")
+        ack_success(
+            self.logger,
+            "create_span",
+            pipeline_run_id=pipeline_run_id,
+            span_uid=res.uid,
+            span_id=res.id,
+        )
+        return res
+
+    def get_child_spans(self, span_id: int) -> list[Span]:
+        url = self._client.pipeline_urls.span_child_spans(span_id)
+        response = self._client._get(url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from child spans, got {type(response).__name__}")
+        raw = response.get("childSpans") or []
+        if not isinstance(raw, list):
+            raise TypeError("childSpans must be a list when present")
+        out: list[Span] = []
+        for item in raw:
+            if not isinstance(item, dict):
+                raise TypeError("Each child span must be an object")
+            s = span_from_api_dict(item)
+            if s is not None:
+                out.append(s)
+        ack_success(self.logger, "get_child_spans", span_id=span_id, count=len(out))
+        return out
+
+    def create_span_event(self, request: SpanEventRequest) -> SpanEvent:
+        if request.event.span_id is None:
+            raise ValueError("To create a span event, spanId is required.")
+        url = self._client.pipeline_urls.span_events(request.event.span_id)
+        payload = _sanitize_for_json(request.to_dict())
+        self.logger.info(f"create span event payload: {payload}")
+        if not isinstance(payload, dict):
+            raise TypeError("Span event payload must serialize to a JSON object")
+        response = self._client._post(url, payload=payload)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from create span event, got {type(response).__name__}")
+        ev = response.get("event")
+        if not isinstance(ev, dict):
+            raise TypeError("Create span event response must include an 'event' object")
+        out = SpanEvent.from_dict(ev)
+        if out is None:
+            raise TypeError("Could not parse span event from response")
+        ack_success(
+            self.logger,
+            "create_span_event",
+            span_id=request.event.span_id,
+            event_uid=request.event.event_uid,
+        )
+        return out
+
+    # get pipeline by uid or id
+    def get_pipeline(self, pipeline_identity: str) -> PipelineResponse:
+        """
+            Description:
+                To get an existing pipeline from the catalog
+        :param pipeline_identity: uid or id of the pipeline
+        :return:(Pipeline) pipeline class instance
+        """
+        url = self._client.pipeline_urls.pipeline_by_identity(pipeline_identity)
+        response = self._client._get(url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from get pipeline, got {type(response).__name__}")
+        out = pipeline_response_from_api_dict(response)
+        pid = out.pipeline.id if out.pipeline is not None else None
+        ack_success(self.logger, "get_pipeline", pipeline_identity=pipeline_identity, pipeline_id=pid)
+        return out
+
+    def get_root_span(self, pipeline_run_id: int) -> Optional[SpanResource]:
+        from acceldata.models.sdk.pipeline.span_resource import wrap_fetched_span
+
+        url = self._client.pipeline_urls.pipeline_run_spans(pipeline_run_id)
+        response = self._client._get(url, params={"onlyRootSpan": "true"})
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from root span, got {type(response).__name__}")
+        spans = response.get("spans") or []
+        if not spans:
+            # Some deployments omit root when ``onlyRootSpan`` is set; fall back to full listing.
+            full = self._client._get(url)
+            if not isinstance(full, dict):
+                raise TypeError(f"Expected dict from spans listing, got {type(full).__name__}")
+            for item in full.get("spans") or []:
+                if not isinstance(item, dict):
+                    raise TypeError("Each span must be an object")
+                if item.get("parentSpanId") is not None:
+                    continue
+                span = span_from_api_dict(item)
+                if span is not None:
+                    return wrap_fetched_span(self, span)
+            return None
+        first = spans[0]
+        if not isinstance(first, dict):
+            raise TypeError("spans[0] must be an object")
+        span = span_from_api_dict(first)
+        if span is None:
+            raise TypeError("Could not parse root span from response")
+        return wrap_fetched_span(self, span)
+
+    def get_pipelines(self) -> PipelinesListingResponse:
+        """
+            Description:
+                To get all pipelines from ADOC
+        :return:(List[PipelineListingInfo]) PipelineListingInfo class instance
+        """
+        url = self._client.pipeline_urls.pipelines_summary()
+        response = self._client._get(url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from list pipelines, got {type(response).__name__}")
+        listing = PipelinesListingResponse.from_dict(response)
+        ack_success(self.logger, "get_pipelines")
+        return listing
+
+    def get_run_details(self, pipeline_id: int, version_id: int) -> PipelineDetailsResponse:
+        url = self._client.pipeline_urls.pipeline_graph_by_version_id(pipeline_id, version_id)
+        response = self._client._get(url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from get_run_details, got {type(response).__name__}")
+        details = PipelineDetailsResponse.from_dict(response)
+        ack_success(self.logger, "get_run_details", pipeline_id=pipeline_id, version_id=version_id)
+        return details
+
+    def get_spans(self, pipeline_run_id: int) -> SpansResponse:
+        """
+            Description:
+                List spans for the pipeline run.
+        :return: :class:`~acceldata.models.api.pipeline.spans_response.SpansResponse`
+        """
+        url = self._client.pipeline_urls.pipeline_run_spans(pipeline_run_id=pipeline_run_id)
+        response = self._client._get(url=url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from get_spans, got {type(response).__name__}")
+        spans = spans_response_from_api_dict(response)
+        ack_success(self.logger, "get_spans", pipeline_run_id=pipeline_run_id)
+        return spans
+
+    def get_span(self, pipeline_run_id: int, span_identity: str) -> Span:
+        url = self._client.pipeline_urls.pipeline_run_span_by_uid(pipeline_run_id=pipeline_run_id,
+                                                                  span_identity=span_identity)
+        response = self._client._get(url=url)
+        if not isinstance(response, dict):
+            raise TypeError(f"Expected dict from get_span, got {type(response).__name__}")
+        # API returns SpanRequest: { "span": { ... }, "associatedJobUids": [...] }
+        if isinstance(response.get("span"), dict):
+            envelope = span_request_from_api_dict(response)
+            if envelope is None or envelope.span is None:
+                raise TypeError("get_span response missing span")
+            span = envelope.span
+        else:
+            span = span_from_api_dict(response)
+        ack_success(
+            self.logger,
+            "get_span",
+            pipeline_run_id=pipeline_run_id,
+            span_identity=span_identity,
+            span_id=span.id,
+        )
+        return span