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

acceldata/models/sdk/pipeline/pipeline_run_resource.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+from acceldata.models.api.pipeline.pipeline_details_response import PipelineDetailsResponse
+from acceldata.models.api.pipeline.node import Node
+from acceldata.models.api.pipeline.pipeline_run import PipelineRun
+from acceldata.models.api.pipeline.spans_response import SpansResponse
+from acceldata.models.sdk.pipeline.create_job_input import CreateJobInput, build_create_job_payload
+from acceldata.models.sdk.pipeline.pipeline_run_input import (
+    PipelineRunInput,
+    PipelineRunResult,
+    PipelineRunStatus,
+)
+from acceldata.utils.span_helpers import ensure_span_for_job, span_create_payload
+from acceldata.models.sdk.pipeline.span_resource import SpanResource, wrap_fetched_span
+from acceldata.services.pipeline_service import PipelineService
+
+logger = logging.getLogger("pipeline_run_resource")
+
+
+class PipelineRunResource:
+    """
+    Customer-facing pipeline run operations that wrap :class:`PipelineRun` without modifying it.
+
+    The wrapped :class:`PipelineRun` is :attr:`pipeline_run`; use :meth:`to_dict` for a wire view.
+    """
+
+    def __init__(self, model: PipelineRun, service: PipelineService) -> None:
+        self._model = model
+        self._service = service
+
+    @property
+    def pipeline_run(self) -> PipelineRun:
+        return self._model
+
+    def to_dict(self) -> Dict[str, Any]:
+        return self._model.to_dict()
+
+    @property
+    def id(self) -> Optional[int]:
+        return self._model.id
+
+    def __repr__(self) -> str:
+        return f"PipelineRunResource(model={self._model!r})"
+
+    def _require_pipeline_id(self) -> int:
+        if self._model.pipeline_id is None:
+            raise ValueError("Pipeline id is required for this operation")
+        return int(self._model.pipeline_id)
+
+    def _require_pipeline_run_id(self) -> int:
+        if self._model.id is None:
+            raise ValueError("Pipeline run id is required for this operation")
+        return int(self._model.id)
+
+    def _require_version_id(self) -> int:
+        if self._model.version_id is None:
+            raise ValueError("Pipeline run version id is required for this operation")
+        return int(self._model.version_id)
+
+    def update_pipeline_run(
+            self,
+            request: PipelineRunInput | None = None,
+            *,
+            context_data: Optional[dict[str, Any]] = None,
+            continuation_id: Optional[str] = None,
+            started_at: Optional[datetime] = None,
+            finished_at: Optional[datetime] = None,
+            result: Optional[PipelineRunResult] = None,
+            status: Optional[PipelineRunStatus] = None,
+    ) -> PipelineRunResource:
+        base = request or PipelineRunInput()
+        eff_context = context_data if context_data is not None else base.context_data
+        eff_continuation = continuation_id if continuation_id is not None else base.continuation_id
+        eff_started = started_at if started_at is not None else base.started_at
+        eff_finished = finished_at if finished_at is not None else base.finished_at
+        eff_result = (
+            result
+            if result is not None
+            else (base.result if base.result is not None else PipelineRunResult.SUCCESS)
+        )
+        eff_status = (
+            status
+            if status is not None
+            else (base.status if base.status is not None else PipelineRunStatus.COMPLETED)
+        )
+        vid = self._model.version_id
+        run = self._service.update_pipeline_run(
+            self._require_pipeline_run_id(),
+            pipeline_id=self._require_pipeline_id(),
+            result=eff_result,
+            status=eff_status,
+            version_id=int(vid) if vid is not None else None,
+            context_data=eff_context,
+            continuation_id=eff_continuation,
+            started_at=eff_started,
+            finished_at=eff_finished,
+        )
+        self._model = run
+        return self
+
+    def create_job(self, job: CreateJobInput) -> Node:
+        if not job.uid or not job.name:
+            raise ValueError(
+                "To create a job, uid and name are required"
+            )
+        effective_run_id = job.pipeline_run_id
+        if effective_run_id is None:
+            effective_run_id = self._require_pipeline_run_id()
+        elif effective_run_id != self._require_pipeline_run_id():
+            raise ValueError(
+                "job.pipeline_run_id must match this pipeline run when set"
+            )
+        job_for_payload = replace(job, pipeline_run_id=effective_run_id)
+        pipeline_id = self._require_pipeline_id()
+        vid = self._model.version_id
+        if vid is None:
+            resolved = self._service.get_pipeline_run(pipeline_run_id=effective_run_id)
+            vid = resolved.version_id
+        payload = build_create_job_payload(
+            job_for_payload,
+            version_id=int(vid) if vid is not None else None,
+        )
+        node_response = self._service.create_job(payload, pipeline_id)
+
+        if job.bounded_by_span:
+            if job.span_uid is None:
+                raise ValueError(
+                    "When bounded_by_span is True, span_uid is required"
+                )
+            ensure_span_for_job(
+                self._service,
+                pipeline_run_id=effective_run_id,
+                span_uid=job.span_uid,
+                job_uid=job.uid,
+                created_at=job.span_created_at,
+                updated_at=job.span_updated_at,
+            )
+        return node_response.node
+
+    # Creates the root span which belongs to the whole pipeline run tracking.
+    def create_root_span(
+            self,
+            uid: str,
+            *,
+            associated_job_uids: Optional[list[str]] = None,
+            with_explicit_time: bool = False,
+            created_at: Optional[datetime] = None,
+            updated_at: Optional[datetime] = None,
+    ) -> SpanResource:
+        """
+        Create a span for this pipeline run (root span when ``parent_span_id`` is ``None``).
+
+        Returns the created :class:`~acceldata.models.sdk.pipeline.span_resource.SpanResource`.
+
+        ``with_explicit_time`` is accepted for compatibility with the earlier Python SDK; span timing is
+        server-defined by the API.
+
+        Pass ``created_at`` / ``updated_at`` to set span record timestamps on create (earlier SDK
+        ``createdAt`` on the span payload); use with explicit-time span lifecycle control.
+        """
+        if not uid:
+            raise ValueError("span uid is required")
+        pr_id = self._require_pipeline_run_id()
+        jobs = list(associated_job_uids or [])
+        res = self._service.create_span(
+            pr_id,
+            span_create_payload(
+                pipeline_run_id=pr_id,
+                uid=uid,
+                associated_job_uids=jobs,
+                created_at=created_at,
+                updated_at=updated_at,
+            ),
+        )
+
+        return SpanResource(
+            pipeline_service=self._service,
+            span=res,
+            new_span=True,
+            with_explicit_time=with_explicit_time,
+        )
+
+    def get_root_span(self) -> Optional[SpanResource]:
+        """Return the root :class:`~acceldata.models.sdk.pipeline.span_resource.SpanResource`, if any."""
+        return self._service.get_root_span(self._require_pipeline_run_id())
+
+    def get_details(self) -> PipelineDetailsResponse:
+        """
+            Description:
+                Get pipeline graph details for this run (nodes, edges).
+        :return: :class:`~acceldata.models.api.pipeline.pipeline_details_response.PipelineDetailsResponse`
+        """
+        response = self._service.get_run_details(pipeline_id=self._require_pipeline_id(),
+                                                 version_id=self._require_version_id())
+        return response
+
+    def get_spans(self) -> SpansResponse:
+        response = self._service.get_spans(self._require_pipeline_run_id())
+        return response
+
+    def get_span(self, span_identity: str) -> SpanResource:
+        s = self._service.get_span(self._require_pipeline_run_id(), span_identity=span_identity)
+        return wrap_fetched_span(self._service, s)

acceldata/models/sdk/pipeline/pipeline_source_type.py

@@ -0,0 +1,20 @@
+# coding: utf-8
+"""``Pipeline.sourceType`` values (aligned with the generated :class:`~acceldata.models.api.pipeline.pipeline.Pipeline` schema)."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class PipelineSourceType(str, Enum):
+    """How the pipeline is orchestrated in the data plane (API wire values)."""
+
+    AIRFLOW = "AIRFLOW"
+    AZURE_DATA_FACTORY = "AZURE_DATA_FACTORY"
+    TRINO = "TRINO"
+    OTHERS = "OTHERS"
+    AUTOSYS = "AUTOSYS"
+    SPARK = "SPARK"
+    DBT = "DBT"
+    FIVETRAN = "FIVETRAN"
+    SNAPLOGIC = "SNAPLOGIC"

acceldata/models/sdk/pipeline/span_resource.py

@@ -0,0 +1,327 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from acceldata.models.api.pipeline.create_span_event import CreateSpanEvent
+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.utils.span_helpers import (
+    normalize_span_event_created_at,
+    span_create_payload,
+)
+from acceldata.utils.operation_ack import ack_failure
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from acceldata.services.pipeline_service import PipelineService
+
+
+class SpanEventType(str, Enum):
+    """Pipeline span event ``eventUid`` values (control plane)."""
+
+    SPAN_START = "SPAN.START"
+    SPAN_END = "SPAN.END"
+    SPAN_ABORTED = "SPAN.ABORTED"
+    SPAN_FAILED = "SPAN.FAILED"
+    LOG = "LOG"
+
+
+class SpanContextStatus(str, Enum):
+    """Lifecycle state strings on :class:`~acceldata.models.api.pipeline.span.Span`."""
+
+    INITIALIZED = "INITIALIZED"
+    STARTED = "STARTED"
+    ABORTED = "ABORTED"
+    FINISHED = "FINISHED"
+    FAILED = "FAILED"
+
+
+_TERMINAL_SPAN_STATUSES = frozenset(
+    {
+        "ABORTED",
+        "FINISHED",
+        "SKIPPED",
+        "ERRORED",
+    }
+)
+
+
+class SpanResource:
+    """
+    Resource wrapper for pipeline span lifecycle.
+
+    Emits span lifecycle events and creates child spans for a pipeline run. Uses generated API models
+    (:class:`~acceldata.models.api.pipeline.span.Span`,
+    :class:`~acceldata.models.api.pipeline.span_event.SpanEvent`) and delegates HTTP
+    to :class:`~acceldata.services.pipeline_service.PipelineService`.
+
+    Construction is side-effect free: creating this resource does not emit ``SPAN.START``.
+    Call :meth:`start` explicitly when the span should begin.
+    """
+
+    def __init__(
+        self,
+        *,
+        pipeline_service: PipelineService,
+        span: Span,
+        parent: Optional[SpanResource] = None,
+        new_span: Optional[bool] = None,
+        with_explicit_time: bool = False,
+    ) -> None:
+        if span is None:
+            raise ValueError("span is required")
+        self._pipeline_service = pipeline_service
+        self._model = span
+        self.parent = parent
+        self.children: list[SpanResource] = []
+
+        if new_span is None:
+            new_span = False
+        self.new_span = new_span
+
+        logger.debug(
+            "SpanResource init with_explicit_time=%s (no implicit start)",
+            with_explicit_time,
+        )
+
+    @property
+    def span(self) -> Span:
+        return self._model
+
+    def __repr__(self) -> str:
+        return f"SpanResource(model={self._model!r})"
+
+    def to_dict(self) -> dict[str, Any]:
+        return self._model.to_dict()
+
+    @property
+    def uid(self) -> str:
+        return self._model.uid
+
+    @property
+    def id(self) -> Optional[int]:
+        return self._model.id
+
+    def _status_upper(self) -> str:
+        return (self._model.status or "").upper()
+
+    def _require_span_id(self) -> int:
+        if self._model.id is None:
+            raise ValueError("span id is required for this operation")
+        return int(self._model.id)
+
+    def _require_pipeline_run_id(self) -> int:
+        if self._model.pipeline_run_id is None:
+            raise ValueError("pipelineRunId is required for this operation")
+        return int(self._model.pipeline_run_id)
+
+    def is_span_finished(self) -> bool:
+        st = self._status_upper()
+        if st in _TERMINAL_SPAN_STATUSES:
+            ack_failure(
+                logger,
+                "span_event",
+                "span_already_terminal",
+                status=st,
+                span_uid=self._model.uid,
+            )
+            return True
+        return False
+
+    def _emit(
+        self,
+        event_uid: str,
+        *,
+        context_data: Optional[dict[str, Any]] = None,
+        created_at: Optional[datetime] = None,
+        log_data: Optional[str] = None,
+    ) -> SpanEvent:
+        """
+        Emit a span lifecycle or custom event through the pipeline service.
+
+        :param event_uid: Event uid, similar to event type (normalized for the API).
+        :param log_data: Log data for the event.
+        :param context_data: Optional context data map.
+        :param created_at: Optional explicit wall-clock time for the span event. If not set,
+            the server applies its default (typically the receive time).
+        :return: :class:`~acceldata.models.api.pipeline.span_event.SpanEvent`
+            for the created event.
+
+        ``span_id`` is taken from the wrapped model; requests use :attr:`_pipeline_service`.
+        """
+        created_at = normalize_span_event_created_at(created_at)
+        svc = self._pipeline_service
+        req = SpanEventRequest(
+            # Use API aliases for compatibility with generated models that may
+            # not allow population by field name (snake_case).
+            event=CreateSpanEvent(
+                eventUid=event_uid,
+                spanId=self._require_span_id(),
+                contextData=context_data,
+                createdAt=created_at,
+                logData=log_data,
+            )
+        )
+        return svc.create_span_event(req)
+
+    def start(
+        self,
+        context_data: Optional[dict[str, Any]] = None,
+        created_at: Optional[datetime] = None,
+    ) -> Optional[SpanEvent]:
+        if self.is_span_finished():
+            return None
+
+        if self._status_upper() == SpanContextStatus.STARTED.value and created_at is None:
+            ack_failure(logger, "span_start", "already_started", span_uid=self._model.uid)
+            return None
+        event_res = self._emit(
+            SpanEventType.SPAN_START.value,
+            context_data=context_data,
+            created_at=created_at,
+        )
+        self._model = self._model.model_copy(update={"status": SpanContextStatus.STARTED.value})
+        return event_res
+
+    def end(
+        self,
+        context_data: Optional[dict[str, Any]] = None,
+        created_at: Optional[datetime] = None,
+    ) -> Optional[SpanEvent]:
+        if self.is_span_finished():
+            return None
+        if self._status_upper() != SpanContextStatus.STARTED.value:
+            ack_failure(logger, "span_end", "not_started", span_uid=self._model.uid)
+            return None
+        event_res = self._emit(
+            SpanEventType.SPAN_END.value,
+            context_data=context_data,
+            created_at=created_at,
+        )
+        self._model = self._model.model_copy(update={"status": SpanContextStatus.FINISHED.value})
+        return event_res
+
+    def abort(
+        self,
+        context_data: Optional[dict[str, Any]] = None,
+        created_at: Optional[datetime] = None,
+    ) -> Optional[SpanEvent]:
+        if self.is_span_finished():
+            return None
+        if self._status_upper() != SpanContextStatus.STARTED.value:
+            ack_failure(logger, "span_abort", "not_started", span_uid=self._model.uid)
+            return None
+        event_res = self._emit(
+            SpanEventType.SPAN_ABORTED.value,
+            context_data=context_data,
+            created_at=created_at,
+        )
+        self._model = self._model.model_copy(update={"status": SpanContextStatus.ABORTED.value})
+        return event_res
+
+    def failed(
+        self,
+        context_data: Optional[dict[str, Any]] = None,
+        created_at: Optional[datetime] = None,
+    ) -> Optional[SpanEvent]:
+        if self.is_span_finished():
+            return None
+        if self._status_upper() != SpanContextStatus.STARTED.value:
+            ack_failure(logger, "span_failed", "not_started", span_uid=self._model.uid)
+            return None
+        event_res = self._emit(
+            SpanEventType.SPAN_FAILED.value,
+            context_data=context_data,
+            created_at=created_at,
+        )
+        # Span status for a failed run uses ERRORED (see Span.status description).
+        self._model = self._model.model_copy(update={"status": "ERRORED"})
+        return event_res
+
+    def send_event(self, span_event: Any) -> Optional[SpanEvent]:
+        """Send a custom event; ``span_event`` must expose ``event_uid`` and optional context/log fields."""
+        if self.is_span_finished():
+            return None
+        if self._status_upper() != SpanContextStatus.STARTED.value:
+            ack_failure(logger, "span_send_event", "not_started", span_uid=self._model.uid)
+            return None
+        event_uid = getattr(span_event, "event_uid", None)
+        if not event_uid:
+            raise ValueError("span_event.event_uid is required")
+        log_data = getattr(span_event, "log_data", None)
+        return self._emit(
+            str(event_uid),
+            context_data=getattr(span_event, "context_data", None),
+            created_at=getattr(span_event, "created_at", None),
+            log_data=log_data,
+        )
+
+    def create_child_span(
+        self,
+        uid: str,
+        associated_job_uids: Optional[List[str]] = None,
+        with_explicit_time: bool = False,
+        created_at: Optional[datetime] = None,
+        updated_at: Optional[datetime] = None,
+    ) -> SpanResource:
+        if not uid:
+            raise ValueError("To create a child span, uid is required")
+        if associated_job_uids is not None and not isinstance(associated_job_uids, list):
+            raise TypeError("associated_job_uids must be a list")
+        jobs = list(associated_job_uids or [])
+        svc = self._pipeline_service
+        pr_id = self._require_pipeline_run_id()
+        parent_id = self._require_span_id()
+        payload = span_create_payload(
+            pipeline_run_id=pr_id,
+            uid=uid,
+            associated_job_uids=jobs,
+            parent_span_id=parent_id,
+            created_at=created_at,
+            updated_at=updated_at,
+        )
+        res = svc.create_span(pr_id, payload)
+        child = SpanResource(
+            pipeline_service=svc,
+            span=res,
+            parent=self,
+            new_span=True,
+        )
+        logger.debug("create_child_span with_explicit_time=%s", with_explicit_time)
+        self.children.append(child)
+        return child
+
+    def has_children(self) -> bool:
+        return len(self.get_child_spans()) > 0
+
+    def get_child_spans(self) -> list[SpanResource]:
+        raw = self._pipeline_service.get_child_spans(self._require_span_id())
+        return [wrap_fetched_span(self._pipeline_service, s, parent=self) for s in raw]
+
+    def is_root(self) -> bool:
+        return self.parent is None
+
+
+def wrap_fetched_span(
+    pipeline_service: PipelineService,
+    span: Span,
+    *,
+    parent: Optional[SpanResource] = None,
+) -> SpanResource:
+    """
+    Build a :class:`SpanResource` for a span returned from GET APIs (root / by uid).
+
+    Does not emit lifecycle events on construction; use :meth:`~SpanResource.start` /
+    :meth:`~SpanResource.end` as needed.
+    """
+    return SpanResource(
+        pipeline_service=pipeline_service,
+        span=span,
+        parent=parent,
+        new_span=False,
+    )