hatchet-sdk 0.47.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

hatchet_sdk/clients/events.py

@@ -1,13 +1,14 @@
 import asyncio
 import datetime
 import json
-from typing import Any, Dict, List, Optional, TypedDict
-from uuid import uuid4
+from typing import List, cast
 
 import grpc
 from google.protobuf import timestamp_pb2
+from pydantic import BaseModel, Field
 
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.contracts.events_pb2 import (
     BulkPushEventRequest,
     Event,
@@ -16,20 +17,18 @@ 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.metadata import get_metadata
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 
-from ..loader import ClientConfig
-from ..metadata import get_metadata
 
-
-def new_event(conn, config: ClientConfig):
+def new_event(conn: grpc.Channel, config: ClientConfig) -> "EventClient":
     return EventClient(
-        client=EventsServiceStub(conn),
+        client=EventsServiceStub(conn),  # type: ignore[no-untyped-call]
         config=config,
     )
 
 
-def proto_timestamp_now():
+def proto_timestamp_now() -> timestamp_pb2.Timestamp:
     t = datetime.datetime.now().timestamp()
     seconds = int(t)
     nanos = int(t % 1 * 1e9)
@@ -37,19 +36,19 @@ def proto_timestamp_now():
     return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
 
 
-class PushEventOptions(TypedDict, total=False):
-    additional_metadata: Dict[str, str] | None = None
+class PushEventOptions(BaseModel):
+    additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
     namespace: str | None = None
 
 
-class BulkPushEventOptions(TypedDict, total=False):
+class BulkPushEventOptions(BaseModel):
     namespace: str | None = None
 
 
-class BulkPushEventWithMetadata(TypedDict, total=False):
+class BulkPushEventWithMetadata(BaseModel):
     key: str
-    payload: Any
-    additional_metadata: Optional[Dict[str, Any]]  # Optional metadata
+    payload: JSONSerializableMapping = Field(default_factory=dict)
+    additional_metadata: JSONSerializableMapping = Field(default_factory=dict)
 
 
 class EventClient:
@@ -58,126 +57,124 @@ class EventClient:
         self.token = config.token
         self.namespace = config.namespace
 
-    async def async_push(
-        self, event_key, payload, options: Optional[PushEventOptions] = None
+    async def aio_push(
+        self,
+        event_key: str,
+        payload: JSONSerializableMapping,
+        options: PushEventOptions = PushEventOptions(),
     ) -> Event:
         return await asyncio.to_thread(
             self.push, event_key=event_key, payload=payload, options=options
         )
 
-    async def async_bulk_push(
+    async def aio_bulk_push(
         self,
-        events: List[BulkPushEventWithMetadata],
-        options: Optional[BulkPushEventOptions] = None,
+        events: list[BulkPushEventWithMetadata],
+        options: BulkPushEventOptions = BulkPushEventOptions(),
     ) -> List[Event]:
         return await asyncio.to_thread(self.bulk_push, events=events, options=options)
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     @tenacity_retry
-    def push(self, event_key, payload, options: PushEventOptions = None) -> Event:
-        namespace = self.namespace
-
-        if (
-            options is not None
-            and "namespace" in options
-            and options["namespace"] is not None
-        ):
-            namespace = options.pop("namespace")
-
+    def push(
+        self,
+        event_key: str,
+        payload: JSONSerializableMapping,
+        options: PushEventOptions = PushEventOptions(),
+    ) -> Event:
+        namespace = options.namespace or self.namespace
         namespaced_event_key = namespace + event_key
 
         try:
-            meta = dict() if options is None else options["additional_metadata"]
-            meta_bytes = None if meta is None else json.dumps(meta).encode("utf-8")
+            meta = options.additional_metadata
+            meta_bytes = None if meta is None else json.dumps(meta)
         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:
+            payload_str = json.dumps(payload)
+        except (TypeError, ValueError) as e:
             raise ValueError(f"Error encoding payload: {e}")
 
         request = PushEventRequest(
             key=namespaced_event_key,
-            payload=payload_bytes,
+            payload=payload_str,
             eventTimestamp=proto_timestamp_now(),
             additionalMetadata=meta_bytes,
         )
 
-        return self.client.Push(request, metadata=get_metadata(self.token))
+        return cast(Event, self.client.Push(request, metadata=get_metadata(self.token)))
+
+    def _create_push_event_request(
+        self,
+        event: BulkPushEventWithMetadata,
+        namespace: str,
+    ) -> PushEventRequest:
+        event_key = namespace + event.key
+        payload = event.payload
+
+        meta = event.additional_metadata
+
+        try:
+            meta_str = json.dumps(meta)
+        except Exception as e:
+            raise ValueError(f"Error encoding meta: {e}")
+
+        try:
+            serialized_payload = json.dumps(payload)
+        except (TypeError, ValueError) as e:
+            raise ValueError(f"Error serializing payload: {e}")
+
+        return PushEventRequest(
+            key=event_key,
+            payload=serialized_payload,
+            eventTimestamp=proto_timestamp_now(),
+            additionalMetadata=meta_str,
+        )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
     @tenacity_retry
     def bulk_push(
         self,
         events: List[BulkPushEventWithMetadata],
-        options: BulkPushEventOptions = None,
+        options: BulkPushEventOptions = BulkPushEventOptions(),
     ) -> List[Event]:
-        namespace = self.namespace
-
-        if (
-            options is not None
-            and "namespace" in options
-            and options["namespace"] is not None
-        ):
-            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)
-
-        bulk_request = BulkPushEventRequest(events=bulk_events)
+        namespace = options.namespace or self.namespace
+
+        bulk_request = BulkPushEventRequest(
+            events=[
+                self._create_push_event_request(event, namespace) for event in events
+            ]
+        )
 
         response = self.client.BulkPush(bulk_request, metadata=get_metadata(self.token))
 
-        return response.events
+        return cast(
+            list[Event],
+            response.events,
+        )
 
-    def log(self, message: str, step_run_id: str):
-        try:
-            request = PutLogRequest(
-                stepRunId=step_run_id,
-                createdAt=proto_timestamp_now(),
-                message=message,
-            )
+    def log(self, message: str, step_run_id: str) -> None:
+        request = PutLogRequest(
+            stepRunId=step_run_id,
+            createdAt=proto_timestamp_now(),
+            message=message,
+        )
 
-            self.client.PutLog(request, metadata=get_metadata(self.token))
-        except Exception as e:
-            raise ValueError(f"Error logging: {e}")
+        self.client.PutLog(request, metadata=get_metadata(self.token))
 
-    def stream(self, data: str | bytes, step_run_id: str):
-        try:
-            if isinstance(data, str):
-                data_bytes = data.encode("utf-8")
-            elif isinstance(data, bytes):
-                data_bytes = data
-            else:
-                raise ValueError("Invalid data type. Expected str, bytes, or file.")
-
-            request = PutStreamEventRequest(
-                stepRunId=step_run_id,
-                createdAt=proto_timestamp_now(),
-                message=data_bytes,
-            )
-            self.client.PutStreamEvent(request, metadata=get_metadata(self.token))
-        except Exception as e:
-            raise ValueError(f"Error putting stream event: {e}")
+    def stream(self, data: str | bytes, step_run_id: str) -> None:
+        if isinstance(data, str):
+            data_bytes = data.encode("utf-8")
+        elif isinstance(data, bytes):
+            data_bytes = data
+        else:
+            raise ValueError("Invalid data type. Expected str, bytes, or file.")
+
+        request = PutStreamEventRequest(
+            stepRunId=step_run_id,
+            createdAt=proto_timestamp_now(),
+            message=data_bytes,
+        )
+
+        self.client.PutStreamEvent(request, metadata=get_metadata(self.token))

hatchet_sdk/clients/rest/__init__.py

@@ -28,11 +28,13 @@ from hatchet_sdk.clients.rest.api.rate_limits_api import RateLimitsApi
 from hatchet_sdk.clients.rest.api.slack_api import SlackApi
 from hatchet_sdk.clients.rest.api.sns_api import SNSApi
 from hatchet_sdk.clients.rest.api.step_run_api import StepRunApi
+from hatchet_sdk.clients.rest.api.task_api import TaskApi
 from hatchet_sdk.clients.rest.api.tenant_api import TenantApi
 from hatchet_sdk.clients.rest.api.user_api import UserApi
 from hatchet_sdk.clients.rest.api.worker_api import WorkerApi
 from hatchet_sdk.clients.rest.api.workflow_api import WorkflowApi
 from hatchet_sdk.clients.rest.api.workflow_run_api import WorkflowRunApi
+from hatchet_sdk.clients.rest.api.workflow_runs_api import WorkflowRunsApi
 from hatchet_sdk.clients.rest.api_client import ApiClient
 
 # import ApiClient
@@ -204,6 +206,7 @@ from hatchet_sdk.clients.rest.models.tenant_resource_policy import TenantResourc
 from hatchet_sdk.clients.rest.models.tenant_step_run_queue_metrics import (
     TenantStepRunQueueMetrics,
 )
+from hatchet_sdk.clients.rest.models.tenant_version import TenantVersion
 from hatchet_sdk.clients.rest.models.trigger_workflow_run_request import (
     TriggerWorkflowRunRequest,
 )
@@ -225,6 +228,24 @@ from hatchet_sdk.clients.rest.models.user_tenant_memberships_list import (
     UserTenantMembershipsList,
 )
 from hatchet_sdk.clients.rest.models.user_tenant_public import UserTenantPublic
+from hatchet_sdk.clients.rest.models.v1_cancel_task_request import V1CancelTaskRequest
+from hatchet_sdk.clients.rest.models.v1_dag_children import V1DagChildren
+from hatchet_sdk.clients.rest.models.v1_replay_task_request import V1ReplayTaskRequest
+from hatchet_sdk.clients.rest.models.v1_task import V1Task
+from hatchet_sdk.clients.rest.models.v1_task_event import V1TaskEvent
+from hatchet_sdk.clients.rest.models.v1_task_event_list import V1TaskEventList
+from hatchet_sdk.clients.rest.models.v1_task_event_type import V1TaskEventType
+from hatchet_sdk.clients.rest.models.v1_task_filter import V1TaskFilter
+from hatchet_sdk.clients.rest.models.v1_task_point_metric import V1TaskPointMetric
+from hatchet_sdk.clients.rest.models.v1_task_point_metrics import V1TaskPointMetrics
+from hatchet_sdk.clients.rest.models.v1_task_run_metric import V1TaskRunMetric
+from hatchet_sdk.clients.rest.models.v1_task_run_status import V1TaskRunStatus
+from hatchet_sdk.clients.rest.models.v1_task_status import V1TaskStatus
+from hatchet_sdk.clients.rest.models.v1_task_summary import V1TaskSummary
+from hatchet_sdk.clients.rest.models.v1_task_summary_list import V1TaskSummaryList
+from hatchet_sdk.clients.rest.models.v1_workflow_run import V1WorkflowRun
+from hatchet_sdk.clients.rest.models.v1_workflow_run_details import V1WorkflowRunDetails
+from hatchet_sdk.clients.rest.models.v1_workflow_type import V1WorkflowType
 from hatchet_sdk.clients.rest.models.webhook_worker import WebhookWorker
 from hatchet_sdk.clients.rest.models.webhook_worker_create_request import (
     WebhookWorkerCreateRequest,
@@ -263,6 +284,9 @@ from hatchet_sdk.clients.rest.models.workflow_run_order_by_field import (
     WorkflowRunOrderByField,
 )
 from hatchet_sdk.clients.rest.models.workflow_run_shape import WorkflowRunShape
+from hatchet_sdk.clients.rest.models.workflow_run_shape_item_for_workflow_run_details import (
+    WorkflowRunShapeItemForWorkflowRunDetails,
+)
 from hatchet_sdk.clients.rest.models.workflow_run_status import WorkflowRunStatus
 from hatchet_sdk.clients.rest.models.workflow_run_triggered_by import (
     WorkflowRunTriggeredBy,

hatchet_sdk/clients/rest/api/__init__.py

@@ -12,8 +12,10 @@ from hatchet_sdk.clients.rest.api.rate_limits_api import RateLimitsApi
 from hatchet_sdk.clients.rest.api.slack_api import SlackApi
 from hatchet_sdk.clients.rest.api.sns_api import SNSApi
 from hatchet_sdk.clients.rest.api.step_run_api import StepRunApi
+from hatchet_sdk.clients.rest.api.task_api import TaskApi
 from hatchet_sdk.clients.rest.api.tenant_api import TenantApi
 from hatchet_sdk.clients.rest.api.user_api import UserApi
 from hatchet_sdk.clients.rest.api.worker_api import WorkerApi
 from hatchet_sdk.clients.rest.api.workflow_api import WorkflowApi
 from hatchet_sdk.clients.rest.api.workflow_run_api import WorkflowRunApi
+from hatchet_sdk.clients.rest.api.workflow_runs_api import WorkflowRunsApi