hatchet-sdk 0.40.0a9__py3-none-any.whl → 0.41.0__py3-none-any.whl

hatchet_sdk/clients/admin.py

@@ -23,6 +23,13 @@ from hatchet_sdk.contracts.workflows_pb2 import (
     WorkflowVersion,
 )
 from hatchet_sdk.contracts.workflows_pb2_grpc import WorkflowServiceStub
+from hatchet_sdk.utils.serialization import flatten
+from hatchet_sdk.utils.tracing import (
+    create_carrier,
+    create_tracer,
+    inject_carrier_into_metadata,
+    parse_carrier_from_metadata,
+)
 from hatchet_sdk.workflow_run import RunRef, WorkflowRunRef
 
 from ..loader import ClientConfig
@@ -164,6 +171,7 @@ class AdminClientAioImpl(AdminClientBase):
         self.token = config.token
         self.listener_client = new_listener(config)
         self.namespace = config.namespace
+        self.otel_tracer = create_tracer(config=config)
 
     async def run(
         self,
@@ -186,44 +194,74 @@ class AdminClientAioImpl(AdminClientBase):
     async def run_workflow(
         self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
     ) -> WorkflowRunRef:
-        try:
-            if not self.pooled_workflow_listener:
-                self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+        ctx = parse_carrier_from_metadata(
+            (options or {}).get("additional_metadata", {})
+        )
 
-            namespace = self.namespace
+        with self.otel_tracer.start_as_current_span(
+            f"hatchet.async_run_workflow.{workflow_name}", context=ctx
+        ) as span:
+            carrier = create_carrier()
 
-            if (
-                options is not None
-                and "namespace" in options
-                and options["namespace"] is not None
-            ):
-                namespace = options["namespace"]
-                del options["namespace"]
+            try:
+                if not self.pooled_workflow_listener:
+                    self.pooled_workflow_listener = PooledWorkflowRunListener(
+                        self.config
+                    )
 
-            if namespace != "" and not workflow_name.startswith(self.namespace):
-                workflow_name = f"{namespace}{workflow_name}"
+                namespace = self.namespace
 
-            request = self._prepare_workflow_request(workflow_name, input, options)
-            resp: TriggerWorkflowResponse = await self.aio_client.TriggerWorkflow(
-                request,
-                metadata=get_metadata(self.token),
-            )
-            return WorkflowRunRef(
-                workflow_run_id=resp.workflow_run_id,
-                workflow_listener=self.pooled_workflow_listener,
-                workflow_run_event_listener=self.listener_client,
-            )
-        except grpc.RpcError as e:
-            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
-                raise DedupeViolationErr(e.details())
+                if (
+                    options is not None
+                    and "namespace" in options
+                    and options["namespace"] is not None
+                ):
+                    namespace = options.pop("namespace")
 
-            raise ValueError(f"gRPC error: {e}")
+                if namespace != "" and not workflow_name.startswith(self.namespace):
+                    workflow_name = f"{namespace}{workflow_name}"
+
+                if options is not None and "additional_metadata" in options:
+                    options["additional_metadata"] = inject_carrier_into_metadata(
+                        options["additional_metadata"], carrier
+                    )
+                    span.set_attributes(
+                        flatten(
+                            options["additional_metadata"], parent_key="", separator="."
+                        )
+                    )
+
+                request = self._prepare_workflow_request(workflow_name, input, options)
+
+                span.add_event(
+                    "Triggering workflow", attributes={"workflow_name": workflow_name}
+                )
+
+                resp: TriggerWorkflowResponse = await self.aio_client.TriggerWorkflow(
+                    request,
+                    metadata=get_metadata(self.token),
+                )
+
+                span.add_event(
+                    "Received workflow response",
+                    attributes={"workflow_name": workflow_name},
+                )
+
+                return WorkflowRunRef(
+                    workflow_run_id=resp.workflow_run_id,
+                    workflow_listener=self.pooled_workflow_listener,
+                    workflow_run_event_listener=self.listener_client,
+                )
+            except grpc.RpcError as e:
+                if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                    raise DedupeViolationErr(e.details())
+
+                raise ValueError(f"gRPC error: {e}")
 
     @tenacity_retry
     async def run_workflows(
         self, workflows: List[WorkflowRunDict], options: TriggerWorkflowOptions = None
     ) -> List[WorkflowRunRef]:
-
         if len(workflows) == 0:
             raise ValueError("No workflows to run")
         try:
@@ -243,7 +281,6 @@ class AdminClientAioImpl(AdminClientBase):
             workflow_run_requests: TriggerWorkflowRequest = []
 
             for workflow in workflows:
-
                 workflow_name = workflow["workflow_name"]
                 input_data = workflow["input"]
                 options = workflow["options"]
@@ -360,6 +397,7 @@ class AdminClient(AdminClientBase):
         self.token = config.token
         self.listener_client = new_listener(config)
         self.namespace = config.namespace
+        self.otel_tracer = create_tracer(config=config)
 
     @tenacity_retry
     def put_workflow(
@@ -408,7 +446,6 @@ class AdminClient(AdminClientBase):
         options: ScheduleTriggerWorkflowOptions = None,
     ) -> WorkflowVersion:
         try:
-
             namespace = self.namespace
 
             if (
@@ -436,55 +473,88 @@ class AdminClient(AdminClientBase):
 
             raise ValueError(f"gRPC error: {e}")
 
+    ## TODO: `options` is treated as a dict (wrong type hint)
+    ## TODO: `any` type hint should come from `typing`
     @tenacity_retry
     def run_workflow(
         self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
     ) -> WorkflowRunRef:
-        try:
-            if not self.pooled_workflow_listener:
-                self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+        ctx = parse_carrier_from_metadata(
+            (options or {}).get("additional_metadata", {})
+        )
 
-            namespace = self.namespace
+        with self.otel_tracer.start_as_current_span(
+            f"hatchet.run_workflow.{workflow_name}", context=ctx
+        ) as span:
+            carrier = create_carrier()
 
-            if (
-                options is not None
-                and "namespace" in options
-                and options["namespace"] is not None
-            ):
-                namespace = options["namespace"]
-                del options["namespace"]
+            try:
+                if not self.pooled_workflow_listener:
+                    self.pooled_workflow_listener = PooledWorkflowRunListener(
+                        self.config
+                    )
 
-            if namespace != "" and not workflow_name.startswith(self.namespace):
-                workflow_name = f"{namespace}{workflow_name}"
+                namespace = self.namespace
 
-            request = self._prepare_workflow_request(workflow_name, input, options)
-            resp: TriggerWorkflowResponse = self.client.TriggerWorkflow(
-                request,
-                metadata=get_metadata(self.token),
-            )
-            return WorkflowRunRef(
-                workflow_run_id=resp.workflow_run_id,
-                workflow_listener=self.pooled_workflow_listener,
-                workflow_run_event_listener=self.listener_client,
-            )
-        except grpc.RpcError as e:
-            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
-                raise DedupeViolationErr(e.details())
+                ## TODO: Factor this out - it's repeated a lot of places
+                if (
+                    options is not None
+                    and "namespace" in options
+                    and options["namespace"] is not None
+                ):
+                    namespace = options.pop("namespace")
 
-            raise ValueError(f"gRPC error: {e}")
+                if options is not None and "additional_metadata" in options:
+                    options["additional_metadata"] = inject_carrier_into_metadata(
+                        options["additional_metadata"], carrier
+                    )
+
+                    span.set_attributes(
+                        flatten(
+                            options["additional_metadata"], parent_key="", separator="."
+                        )
+                    )
+
+                if namespace != "" and not workflow_name.startswith(self.namespace):
+                    workflow_name = f"{namespace}{workflow_name}"
+
+                request = self._prepare_workflow_request(workflow_name, input, options)
+
+                span.add_event(
+                    "Triggering workflow", attributes={"workflow_name": workflow_name}
+                )
+
+                resp: TriggerWorkflowResponse = self.client.TriggerWorkflow(
+                    request,
+                    metadata=get_metadata(self.token),
+                )
+
+                span.add_event(
+                    "Received workflow response",
+                    attributes={"workflow_name": workflow_name},
+                )
+
+                return WorkflowRunRef(
+                    workflow_run_id=resp.workflow_run_id,
+                    workflow_listener=self.pooled_workflow_listener,
+                    workflow_run_event_listener=self.listener_client,
+                )
+            except grpc.RpcError as e:
+                if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                    raise DedupeViolationErr(e.details())
+
+                raise ValueError(f"gRPC error: {e}")
 
     @tenacity_retry
     def run_workflows(
         self, workflows: List[WorkflowRunDict], options: TriggerWorkflowOptions = None
     ) -> list[WorkflowRunRef]:
-
         workflow_run_requests: TriggerWorkflowRequest = []
         try:
             if not self.pooled_workflow_listener:
                 self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
 
             for workflow in workflows:
-
                 workflow_name = workflow["workflow_name"]
                 input_data = workflow["input"]
                 options = workflow["options"]

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 import time
 from dataclasses import dataclass, field
-from typing import AsyncGenerator, List, Optional
+from typing import Any, AsyncGenerator, List, Optional
 
 import grpc
 from grpc._cython import cygrpc
@@ -23,6 +23,7 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.backoff import exp_backoff_sleep
+from hatchet_sdk.utils.serialization import flatten
 
 from ...loader import ClientConfig
 from ...metadata import get_metadata
@@ -88,6 +89,29 @@ class Action:
         if not isinstance(self.additional_metadata, dict):
             self.additional_metadata = {}
 
+    @property
+    def otel_attributes(self) -> dict[str, Any]:
+        return flatten(
+            xs={
+                "worker_id": self.worker_id,
+                "tenant_id": self.tenant_id,
+                "workflow_run_id": self.workflow_run_id,
+                "get_group_key_run_id": self.get_group_key_run_id,
+                "job_id": self.job_id,
+                "job_name": self.job_name,
+                "job_run_id": self.job_run_id,
+                "step_id": self.step_id,
+                "step_run_id": self.step_run_id,
+                "retry_count": self.retry_count,
+                "child_workflow_index": self.child_workflow_index,
+                "child_workflow_key": self.child_workflow_key,
+                "parent_workflow_run_id": self.parent_workflow_run_id,
+                "action_payload": self.action_payload,
+            },
+            parent_key="",
+            separator=".",
+        )
+
 
 START_STEP_RUN = 0
 CANCEL_STEP_RUN = 1
@@ -202,7 +226,6 @@ class ActionListener:
         listener = None
 
         while not self.stop_signal:
-
             if listener is not None:
                 listener.cancel()
 

hatchet_sdk/clients/events.py

@@ -2,6 +2,7 @@ import asyncio
 import datetime
 import json
 from typing import Any, Dict, List, Optional, TypedDict
+from uuid import uuid4
 
 import grpc
 from google.protobuf import timestamp_pb2
@@ -15,6 +16,13 @@ from hatchet_sdk.contracts.events_pb2 import (
     PutStreamEventRequest,
 )
 from hatchet_sdk.contracts.events_pb2_grpc import EventsServiceStub
+from hatchet_sdk.utils.serialization import flatten
+from hatchet_sdk.utils.tracing import (
+    create_carrier,
+    create_tracer,
+    inject_carrier_into_metadata,
+    parse_carrier_from_metadata,
+)
 
 from ..loader import ClientConfig
 from ..metadata import get_metadata
@@ -55,6 +63,7 @@ class EventClient:
         self.client = client
         self.token = config.token
         self.namespace = config.namespace
+        self.otel_tracer = create_tracer(config=config)
 
     async def async_push(
         self, event_key, payload, options: Optional[PushEventOptions] = None
@@ -72,40 +81,54 @@ class EventClient:
 
     @tenacity_retry
     def push(self, event_key, payload, options: PushEventOptions = None) -> Event:
-        namespace = self.namespace
+        ctx = parse_carrier_from_metadata(
+            (options or {}).get("additional_metadata", {})
+        )
 
-        if (
-            options is not None
-            and "namespace" in options
-            and options["namespace"] is not None
-        ):
-            namespace = options["namespace"]
-            del options["namespace"]
+        with self.otel_tracer.start_as_current_span(
+            "hatchet.push", context=ctx
+        ) as span:
+            carrier = create_carrier()
+            namespace = self.namespace
 
-        namespaced_event_key = namespace + event_key
+            if (
+                options is not None
+                and "namespace" in options
+                and options["namespace"] is not None
+            ):
+                namespace = options.pop("namespace")
 
-        try:
-            meta = None if options is None else options["additional_metadata"]
-            meta_bytes = None if meta is None else json.dumps(meta).encode("utf-8")
-        except Exception as e:
-            raise ValueError(f"Error encoding meta: {e}")
+            namespaced_event_key = namespace + event_key
 
-        try:
-            payload_bytes = json.dumps(payload).encode("utf-8")
-        except json.UnicodeEncodeError as e:
-            raise ValueError(f"Error encoding payload: {e}")
-
-        request = PushEventRequest(
-            key=namespaced_event_key,
-            payload=payload_bytes,
-            eventTimestamp=proto_timestamp_now(),
-            additionalMetadata=meta_bytes,
-        )
+            try:
+                meta = inject_carrier_into_metadata(
+                    dict() if options is None else options["additional_metadata"],
+                    carrier,
+                )
+                meta_bytes = None if meta is None else json.dumps(meta).encode("utf-8")
+            except Exception as e:
+                raise ValueError(f"Error encoding meta: {e}")
 
-        try:
-            return self.client.Push(request, metadata=get_metadata(self.token))
-        except grpc.RpcError as e:
-            raise ValueError(f"gRPC error: {e}")
+            span.set_attributes(flatten(meta, parent_key="", separator="."))
+
+            try:
+                payload_bytes = json.dumps(payload).encode("utf-8")
+            except json.UnicodeEncodeError as e:
+                raise ValueError(f"Error encoding payload: {e}")
+
+            request = PushEventRequest(
+                key=namespaced_event_key,
+                payload=payload_bytes,
+                eventTimestamp=proto_timestamp_now(),
+                additionalMetadata=meta_bytes,
+            )
+
+            span.add_event("Pushing event", attributes={"key": namespaced_event_key})
+
+            try:
+                return self.client.Push(request, metadata=get_metadata(self.token))
+            except grpc.RpcError as e:
+                raise ValueError(f"gRPC error: {e}")
 
     @tenacity_retry
     def bulk_push(
@@ -114,41 +137,57 @@ class EventClient:
         options: BulkPushEventOptions = None,
     ) -> List[Event]:
         namespace = self.namespace
+        bulk_push_correlation_id = uuid4()
+        ctx = parse_carrier_from_metadata(
+            (options or {}).get("additional_metadata", {})
+        )
 
         if (
             options is not None
             and "namespace" in options
             and options["namespace"] is not None
         ):
-            namespace = options["namespace"]
-            del options["namespace"]
+            namespace = options.pop("namespace")
 
         bulk_events = []
         for event in events:
-            event_key = namespace + event["key"]
-            payload = event["payload"]
-
-            try:
-                meta = event.get("additional_metadata")
-                meta_bytes = json.dumps(meta).encode("utf-8") if meta else None
-            except Exception as e:
-                raise ValueError(f"Error encoding meta: {e}")
-
-            try:
-                payload_bytes = json.dumps(payload).encode("utf-8")
-            except json.UnicodeEncodeError as e:
-                raise ValueError(f"Error encoding payload: {e}")
-
-            request = PushEventRequest(
-                key=event_key,
-                payload=payload_bytes,
-                eventTimestamp=proto_timestamp_now(),
-                additionalMetadata=meta_bytes,
-            )
-            bulk_events.append(request)
+            with self.otel_tracer.start_as_current_span(
+                "hatchet.bulk_push", context=ctx
+            ) as span:
+                carrier = create_carrier()
+                span.set_attribute(
+                    "bulk_push_correlation_id", str(bulk_push_correlation_id)
+                )
+
+                event_key = namespace + event["key"]
+                payload = event["payload"]
+
+                try:
+                    meta = inject_carrier_into_metadata(
+                        event.get("additional_metadata", {}), carrier
+                    )
+                    meta_bytes = json.dumps(meta).encode("utf-8") if meta else None
+                except Exception as e:
+                    raise ValueError(f"Error encoding meta: {e}")
+
+                span.set_attributes(flatten(meta, parent_key="", separator="."))
+
+                try:
+                    payload_bytes = json.dumps(payload).encode("utf-8")
+                except json.UnicodeEncodeError as e:
+                    raise ValueError(f"Error encoding payload: {e}")
+
+                request = PushEventRequest(
+                    key=event_key,
+                    payload=payload_bytes,
+                    eventTimestamp=proto_timestamp_now(),
+                    additionalMetadata=meta_bytes,
+                )
+                bulk_events.append(request)
 
         bulk_request = BulkPushEventRequest(events=bulk_events)
 
+        span.add_event("Pushing bulk events")
         try:
             response = self.client.BulkPush(
                 bulk_request, metadata=get_metadata(self.token)

hatchet_sdk/clients/rest/__init__.py

@@ -70,6 +70,9 @@ from hatchet_sdk.clients.rest.models.create_api_token_request import (
 from hatchet_sdk.clients.rest.models.create_api_token_response import (
     CreateAPITokenResponse,
 )
+from hatchet_sdk.clients.rest.models.create_cron_workflow_trigger_request import (
+    CreateCronWorkflowTriggerRequest,
+)
 from hatchet_sdk.clients.rest.models.create_event_request import CreateEventRequest
 from hatchet_sdk.clients.rest.models.create_pull_request_from_step_run import (
     CreatePullRequestFromStepRun,
@@ -84,6 +87,11 @@ from hatchet_sdk.clients.rest.models.create_tenant_invite_request import (
     CreateTenantInviteRequest,
 )
 from hatchet_sdk.clients.rest.models.create_tenant_request import CreateTenantRequest
+from hatchet_sdk.clients.rest.models.cron_workflows import CronWorkflows
+from hatchet_sdk.clients.rest.models.cron_workflows_list import CronWorkflowsList
+from hatchet_sdk.clients.rest.models.cron_workflows_order_by_field import (
+    CronWorkflowsOrderByField,
+)
 from hatchet_sdk.clients.rest.models.event import Event
 from hatchet_sdk.clients.rest.models.event_data import EventData
 from hatchet_sdk.clients.rest.models.event_key_list import EventKeyList
@@ -141,6 +149,17 @@ from hatchet_sdk.clients.rest.models.replay_workflow_runs_response import (
     ReplayWorkflowRunsResponse,
 )
 from hatchet_sdk.clients.rest.models.rerun_step_run_request import RerunStepRunRequest
+from hatchet_sdk.clients.rest.models.schedule_workflow_run_request import (
+    ScheduleWorkflowRunRequest,
+)
+from hatchet_sdk.clients.rest.models.scheduled_run_status import ScheduledRunStatus
+from hatchet_sdk.clients.rest.models.scheduled_workflows import ScheduledWorkflows
+from hatchet_sdk.clients.rest.models.scheduled_workflows_list import (
+    ScheduledWorkflowsList,
+)
+from hatchet_sdk.clients.rest.models.scheduled_workflows_order_by_field import (
+    ScheduledWorkflowsOrderByField,
+)
 from hatchet_sdk.clients.rest.models.semaphore_slots import SemaphoreSlots
 from hatchet_sdk.clients.rest.models.slack_webhook import SlackWebhook
 from hatchet_sdk.clients.rest.models.sns_integration import SNSIntegration
@@ -219,6 +238,8 @@ from hatchet_sdk.clients.rest.models.webhook_worker_request_method import (
 from hatchet_sdk.clients.rest.models.worker import Worker
 from hatchet_sdk.clients.rest.models.worker_label import WorkerLabel
 from hatchet_sdk.clients.rest.models.worker_list import WorkerList
+from hatchet_sdk.clients.rest.models.worker_runtime_info import WorkerRuntimeInfo
+from hatchet_sdk.clients.rest.models.worker_runtime_sdks import WorkerRuntimeSDKs
 from hatchet_sdk.clients.rest.models.workflow import Workflow
 from hatchet_sdk.clients.rest.models.workflow_concurrency import WorkflowConcurrency
 from hatchet_sdk.clients.rest.models.workflow_kind import WorkflowKind

hatchet_sdk/clients/rest/api/api_token_api.py

@@ -298,9 +298,10 @@ class APITokenApi:
             _body_params = create_api_token_request
 
         # set the HTTP header `Accept`
-        _header_params["Accept"] = self.api_client.select_header_accept(
-            ["application/json"]
-        )
+        if "Accept" not in _header_params:
+            _header_params["Accept"] = self.api_client.select_header_accept(
+                ["application/json"]
+            )
 
         # set the HTTP header `Content-Type`
         if _content_type:
@@ -570,9 +571,10 @@ class APITokenApi:
         # process the body parameter
 
         # set the HTTP header `Accept`
-        _header_params["Accept"] = self.api_client.select_header_accept(
-            ["application/json"]
-        )
+        if "Accept" not in _header_params:
+            _header_params["Accept"] = self.api_client.select_header_accept(
+                ["application/json"]
+            )
 
         # authentication setting
         _auth_settings: List[str] = ["cookieAuth", "bearerAuth"]
@@ -832,9 +834,10 @@ class APITokenApi:
         # process the body parameter
 
         # set the HTTP header `Accept`
-        _header_params["Accept"] = self.api_client.select_header_accept(
-            ["application/json"]
-        )
+        if "Accept" not in _header_params:
+            _header_params["Accept"] = self.api_client.select_header_accept(
+                ["application/json"]
+            )
 
         # authentication setting
         _auth_settings: List[str] = ["cookieAuth", "bearerAuth"]