hatchet-sdk 0.42.3__py3-none-any.whl → 0.42.5__py3-none-any.whl

hatchet_sdk/clients/admin.py

@@ -41,7 +41,7 @@ def new_admin(config: ClientConfig):
     return AdminClient(config)
 
 
-class ScheduleTriggerWorkflowOptions(TypedDict):
+class ScheduleTriggerWorkflowOptions(TypedDict, total=False):
     parent_id: Optional[str]
     parent_step_run_id: Optional[str]
     child_index: Optional[int]
@@ -49,25 +49,25 @@ class ScheduleTriggerWorkflowOptions(TypedDict):
     namespace: Optional[str]
 
 
-class ChildTriggerWorkflowOptions(TypedDict):
+class ChildTriggerWorkflowOptions(TypedDict, total=False):
     additional_metadata: Dict[str, str] | None = None
     sticky: bool | None = None
 
 
-class ChildWorkflowRunDict(TypedDict):
+class ChildWorkflowRunDict(TypedDict, total=False):
     workflow_name: str
     input: Any
     options: ChildTriggerWorkflowOptions
     key: str | None = None
 
 
-class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions, TypedDict):
+class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions, total=False):
     additional_metadata: Dict[str, str] | None = None
     desired_worker_id: str | None = None
     namespace: str | None = None
 
 
-class WorkflowRunDict(TypedDict):
+class WorkflowRunDict(TypedDict, total=False):
     workflow_name: str
     input: Any
     options: TriggerWorkflowOptions | None
@@ -425,7 +425,7 @@ class AdminClient(AdminClientBase):
         self,
         key: str,
         limit: int,
-        duration: RateLimitDuration = RateLimitDuration.SECOND,
+        duration: Union[RateLimitDuration.Value, str] = RateLimitDuration.SECOND,
     ):
         try:
             self.client.PutRateLimit(

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -1,3 +1,5 @@
+from typing import Any, cast
+
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from hatchet_sdk.clients.dispatcher.action_listener import (
@@ -31,7 +33,7 @@ from ...metadata import get_metadata
 DEFAULT_REGISTER_TIMEOUT = 30
 
 
-def new_dispatcher(config: ClientConfig):
+def new_dispatcher(config: ClientConfig) -> "DispatcherClient":
     return DispatcherClient(config=config)
 
 
@@ -40,10 +42,10 @@ class DispatcherClient:
 
     def __init__(self, config: ClientConfig):
         conn = new_conn(config)
-        self.client = DispatcherStub(conn)
+        self.client = DispatcherStub(conn)  # type: ignore[no-untyped-call]
 
         aio_conn = new_conn(config, True)
-        self.aio_client = DispatcherStub(aio_conn)
+        self.aio_client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
         self.token = config.token
         self.config = config
 
@@ -67,7 +69,7 @@ class DispatcherClient:
 
     async def send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ):
+    ) -> Any:
         try:
             return await self._try_send_step_action_event(action, event_type, payload)
         except Exception as e:
@@ -87,7 +89,7 @@ class DispatcherClient:
     @tenacity_retry
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ):
+    ) -> Any:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -103,6 +105,7 @@ class DispatcherClient:
             eventPayload=payload,
         )
 
+        ## TODO: What does this return?
         return await self.aio_client.SendStepActionEvent(
             event,
             metadata=get_metadata(self.token),
@@ -110,7 +113,7 @@ class DispatcherClient:
 
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
-    ):
+    ) -> Any:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -124,19 +127,21 @@ class DispatcherClient:
             eventPayload=payload,
         )
 
+        ## TODO: What does this return?
         return await self.aio_client.SendGroupKeyActionEvent(
             event,
             metadata=get_metadata(self.token),
         )
 
-    def put_overrides_data(self, data: OverridesData):
-        response: ActionEventResponse = self.client.PutOverridesData(
-            data,
-            metadata=get_metadata(self.token),
+    def put_overrides_data(self, data: OverridesData) -> ActionEventResponse:
+        return cast(
+            ActionEventResponse,
+            self.client.PutOverridesData(
+                data,
+                metadata=get_metadata(self.token),
+            ),
         )
 
-        return response
-
     def release_slot(self, step_run_id: str) -> None:
         self.client.ReleaseSlot(
             ReleaseSlotRequest(stepRunId=step_run_id),
@@ -154,7 +159,9 @@ class DispatcherClient:
             metadata=get_metadata(self.token),
         )
 
-    def upsert_worker_labels(self, worker_id: str, labels: dict[str, str | int]):
+    def upsert_worker_labels(
+        self, worker_id: str | None, labels: dict[str, str | int]
+    ) -> None:
         worker_labels = {}
 
         for key, value in labels.items():
@@ -171,9 +178,9 @@ class DispatcherClient:
 
     async def async_upsert_worker_labels(
         self,
-        worker_id: str,
+        worker_id: str | None,
         labels: dict[str, str | int],
-    ):
+    ) -> None:
         worker_labels = {}
 
         for key, value in labels.items():

hatchet_sdk/clients/events.py

@@ -43,16 +43,16 @@ def proto_timestamp_now():
     return timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
 
 
-class PushEventOptions(TypedDict):
+class PushEventOptions(TypedDict, total=False):
     additional_metadata: Dict[str, str] | None = None
     namespace: str | None = None
 
 
-class BulkPushEventOptions(TypedDict):
+class BulkPushEventOptions(TypedDict, total=False):
     namespace: str | None = None
 
 
-class BulkPushEventWithMetadata(TypedDict):
+class BulkPushEventWithMetadata(TypedDict, total=False):
     key: str
     payload: Any
     additional_metadata: Optional[Dict[str, Any]]  # Optional metadata

hatchet_sdk/clients/rest/api/default_api.py

@@ -48,6 +48,266 @@ class DefaultApi:
             api_client = ApiClient.get_default()
         self.api_client = api_client
 
+    @validate_call
+    async def monitoring_post_run_probe(
+        self,
+        tenant: Annotated[
+            str,
+            Field(
+                min_length=36, strict=True, max_length=36, description="The tenant id"
+            ),
+        ],
+        _request_timeout: Union[
+            None,
+            Annotated[StrictFloat, Field(gt=0)],
+            Tuple[
+                Annotated[StrictFloat, Field(gt=0)], Annotated[StrictFloat, Field(gt=0)]
+            ],
+        ] = None,
+        _request_auth: Optional[Dict[StrictStr, Any]] = None,
+        _content_type: Optional[StrictStr] = None,
+        _headers: Optional[Dict[StrictStr, Any]] = None,
+        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
+    ) -> None:
+        """Detailed Health Probe For the Instance
+
+        Triggers a workflow to check the status of the instance
+
+        :param tenant: The tenant id (required)
+        :type tenant: str
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :type _request_timeout: int, tuple(int, int), optional
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the
+                              authentication in the spec for a single request.
+        :type _request_auth: dict, optional
+        :param _content_type: force content-type for the request.
+        :type _content_type: str, Optional
+        :param _headers: set to override the headers for a single
+                         request; this effectively ignores the headers
+                         in the spec for a single request.
+        :type _headers: dict, optional
+        :param _host_index: set to override the host_index for a single
+                            request; this effectively ignores the host_index
+                            in the spec for a single request.
+        :type _host_index: int, optional
+        :return: Returns the result object.
+        """  # noqa: E501
+
+        _param = self._monitoring_post_run_probe_serialize(
+            tenant=tenant,
+            _request_auth=_request_auth,
+            _content_type=_content_type,
+            _headers=_headers,
+            _host_index=_host_index,
+        )
+
+        _response_types_map: Dict[str, Optional[str]] = {
+            "200": None,
+            "403": "APIErrors",
+        }
+        response_data = await self.api_client.call_api(
+            *_param, _request_timeout=_request_timeout
+        )
+        await response_data.read()
+        return self.api_client.response_deserialize(
+            response_data=response_data,
+            response_types_map=_response_types_map,
+        ).data
+
+    @validate_call
+    async def monitoring_post_run_probe_with_http_info(
+        self,
+        tenant: Annotated[
+            str,
+            Field(
+                min_length=36, strict=True, max_length=36, description="The tenant id"
+            ),
+        ],
+        _request_timeout: Union[
+            None,
+            Annotated[StrictFloat, Field(gt=0)],
+            Tuple[
+                Annotated[StrictFloat, Field(gt=0)], Annotated[StrictFloat, Field(gt=0)]
+            ],
+        ] = None,
+        _request_auth: Optional[Dict[StrictStr, Any]] = None,
+        _content_type: Optional[StrictStr] = None,
+        _headers: Optional[Dict[StrictStr, Any]] = None,
+        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
+    ) -> ApiResponse[None]:
+        """Detailed Health Probe For the Instance
+
+        Triggers a workflow to check the status of the instance
+
+        :param tenant: The tenant id (required)
+        :type tenant: str
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :type _request_timeout: int, tuple(int, int), optional
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the
+                              authentication in the spec for a single request.
+        :type _request_auth: dict, optional
+        :param _content_type: force content-type for the request.
+        :type _content_type: str, Optional
+        :param _headers: set to override the headers for a single
+                         request; this effectively ignores the headers
+                         in the spec for a single request.
+        :type _headers: dict, optional
+        :param _host_index: set to override the host_index for a single
+                            request; this effectively ignores the host_index
+                            in the spec for a single request.
+        :type _host_index: int, optional
+        :return: Returns the result object.
+        """  # noqa: E501
+
+        _param = self._monitoring_post_run_probe_serialize(
+            tenant=tenant,
+            _request_auth=_request_auth,
+            _content_type=_content_type,
+            _headers=_headers,
+            _host_index=_host_index,
+        )
+
+        _response_types_map: Dict[str, Optional[str]] = {
+            "200": None,
+            "403": "APIErrors",
+        }
+        response_data = await self.api_client.call_api(
+            *_param, _request_timeout=_request_timeout
+        )
+        await response_data.read()
+        return self.api_client.response_deserialize(
+            response_data=response_data,
+            response_types_map=_response_types_map,
+        )
+
+    @validate_call
+    async def monitoring_post_run_probe_without_preload_content(
+        self,
+        tenant: Annotated[
+            str,
+            Field(
+                min_length=36, strict=True, max_length=36, description="The tenant id"
+            ),
+        ],
+        _request_timeout: Union[
+            None,
+            Annotated[StrictFloat, Field(gt=0)],
+            Tuple[
+                Annotated[StrictFloat, Field(gt=0)], Annotated[StrictFloat, Field(gt=0)]
+            ],
+        ] = None,
+        _request_auth: Optional[Dict[StrictStr, Any]] = None,
+        _content_type: Optional[StrictStr] = None,
+        _headers: Optional[Dict[StrictStr, Any]] = None,
+        _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
+    ) -> RESTResponseType:
+        """Detailed Health Probe For the Instance
+
+        Triggers a workflow to check the status of the instance
+
+        :param tenant: The tenant id (required)
+        :type tenant: str
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :type _request_timeout: int, tuple(int, int), optional
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the
+                              authentication in the spec for a single request.
+        :type _request_auth: dict, optional
+        :param _content_type: force content-type for the request.
+        :type _content_type: str, Optional
+        :param _headers: set to override the headers for a single
+                         request; this effectively ignores the headers
+                         in the spec for a single request.
+        :type _headers: dict, optional
+        :param _host_index: set to override the host_index for a single
+                            request; this effectively ignores the host_index
+                            in the spec for a single request.
+        :type _host_index: int, optional
+        :return: Returns the result object.
+        """  # noqa: E501
+
+        _param = self._monitoring_post_run_probe_serialize(
+            tenant=tenant,
+            _request_auth=_request_auth,
+            _content_type=_content_type,
+            _headers=_headers,
+            _host_index=_host_index,
+        )
+
+        _response_types_map: Dict[str, Optional[str]] = {
+            "200": None,
+            "403": "APIErrors",
+        }
+        response_data = await self.api_client.call_api(
+            *_param, _request_timeout=_request_timeout
+        )
+        return response_data.response
+
+    def _monitoring_post_run_probe_serialize(
+        self,
+        tenant,
+        _request_auth,
+        _content_type,
+        _headers,
+        _host_index,
+    ) -> RequestSerialized:
+
+        _host = None
+
+        _collection_formats: Dict[str, str] = {}
+
+        _path_params: Dict[str, str] = {}
+        _query_params: List[Tuple[str, str]] = []
+        _header_params: Dict[str, Optional[str]] = _headers or {}
+        _form_params: List[Tuple[str, str]] = []
+        _files: Dict[
+            str, Union[str, bytes, List[str], List[bytes], List[Tuple[str, bytes]]]
+        ] = {}
+        _body_params: Optional[bytes] = None
+
+        # process the path parameters
+        if tenant is not None:
+            _path_params["tenant"] = tenant
+        # process the query parameters
+        # process the header parameters
+        # process the form parameters
+        # process the body parameter
+
+        # set the HTTP header `Accept`
+        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"]
+
+        return self.api_client.param_serialize(
+            method="POST",
+            resource_path="/api/v1/monitoring/{tenant}/probe",
+            path_params=_path_params,
+            query_params=_query_params,
+            header_params=_header_params,
+            body=_body_params,
+            post_params=_form_params,
+            files=_files,
+            auth_settings=_auth_settings,
+            collection_formats=_collection_formats,
+            _host=_host,
+            _request_auth=_request_auth,
+        )
+
     @validate_call
     async def tenant_invite_delete(
         self,

hatchet_sdk/context/context.py

@@ -13,8 +13,8 @@ from hatchet_sdk.clients.rest_client import RestApi
 from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.context.worker_context import WorkerContext
-from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData  # type: ignore
-from hatchet_sdk.contracts.workflows_pb2 import (  # type: ignore[attr-defined]
+from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData
+from hatchet_sdk.contracts.workflows_pb2 import (
     BulkTriggerWorkflowRequest,
     TriggerWorkflowRequest,
 )
@@ -69,7 +69,7 @@ class BaseContext:
             meta = options["additional_metadata"]
 
         ## TODO: Pydantic here to simplify this
-        trigger_options: TriggerWorkflowOptions = {  # type: ignore[typeddict-item]
+        trigger_options: TriggerWorkflowOptions = {
             "parent_id": workflow_run_id,
             "parent_step_run_id": step_run_id,
             "child_key": key,
@@ -149,8 +149,7 @@ class ContextAioImpl(BaseContext):
             key = child_workflow_run.get("key")
             options = child_workflow_run.get("options", {})
 
-            ## TODO: figure out why this is failing
-            trigger_options = self._prepare_workflow_options(key, options, worker_id)  # type: ignore[arg-type]
+            trigger_options = self._prepare_workflow_options(key, options, worker_id)
 
             bulk_trigger_workflow_runs.append(
                 WorkflowRunDict(
@@ -238,14 +237,17 @@ class Context(BaseContext):
             self.input = self.data.get("input", {})
 
     def step_output(self, step: str) -> dict[str, Any] | BaseModel:
-        validators = self.validator_registry.get(step)
+        workflow_validator = next(
+            (v for k, v in self.validator_registry.items() if k.split(":")[-1] == step),
+            None,
+        )
 
         try:
             parent_step_data = cast(dict[str, Any], self.data["parents"][step])
         except KeyError:
             raise ValueError(f"Step output for '{step}' not found")
 
-        if validators and (v := validators.step_output):
+        if workflow_validator and (v := workflow_validator.step_output):
             return v.model_validate(parent_step_data)
 
         return parent_step_data

hatchet_sdk/context/worker_context.py

@@ -2,7 +2,7 @@ from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
 
 
 class WorkerContext:
-    _worker_id: str = None
+    _worker_id: str | None = None
     _registered_workflow_names: list[str] = []
     _labels: dict[str, str | int] = {}
 
@@ -10,18 +10,18 @@ class WorkerContext:
         self._labels = labels
         self.client = client
 
-    def labels(self):
+    def labels(self) -> dict[str, str | int]:
         return self._labels
 
-    def upsert_labels(self, labels: dict[str, str | int]):
+    def upsert_labels(self, labels: dict[str, str | int]) -> None:
         self.client.upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    async def async_upsert_labels(self, labels: dict[str, str | int]):
+    async def async_upsert_labels(self, labels: dict[str, str | int]) -> None:
         await self.client.async_upsert_worker_labels(self._worker_id, labels)
         self._labels.update(labels)
 
-    def id(self) -> str:
+    def id(self) -> str | None:
         return self._worker_id
 
     # def has_workflow(self, workflow_name: str):

hatchet_sdk/contracts/dispatcher_pb2.py

@@ -1,12 +1,22 @@
 # -*- coding: utf-8 -*-
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
 # source: dispatcher.proto
-# Protobuf Python Version: 4.25.1
+# Protobuf Python Version: 5.28.1
 """Generated protocol buffer code."""
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
 from google.protobuf import symbol_database as _symbol_database
 from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    28,
+    1,
+    '',
+    'dispatcher.proto'
+)
 # @@protoc_insertion_point(imports)
 
 _sym_db = _symbol_database.Default()
@@ -20,12 +30,12 @@ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x10\x64ispatcher.
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
 _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'dispatcher_pb2', _globals)
-if _descriptor._USE_C_DESCRIPTORS == False:
-  _globals['DESCRIPTOR']._options = None
+if not _descriptor._USE_C_DESCRIPTORS:
+  _globals['DESCRIPTOR']._loaded_options = None
   _globals['DESCRIPTOR']._serialized_options = b'ZEgithub.com/hatchet-dev/hatchet/internal/services/dispatcher/contracts'
-  _globals['_WORKERREGISTERREQUEST_LABELSENTRY']._options = None
+  _globals['_WORKERREGISTERREQUEST_LABELSENTRY']._loaded_options = None
   _globals['_WORKERREGISTERREQUEST_LABELSENTRY']._serialized_options = b'8\001'
-  _globals['_UPSERTWORKERLABELSREQUEST_LABELSENTRY']._options = None
+  _globals['_UPSERTWORKERLABELSREQUEST_LABELSENTRY']._loaded_options = None
   _globals['_UPSERTWORKERLABELSREQUEST_LABELSENTRY']._serialized_options = b'8\001'
   _globals['_SDKS']._serialized_start=3484
   _globals['_SDKS']._serialized_end=3539