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

hatchet_sdk/connection.py

@@ -1,26 +1,36 @@
 import os
-from typing import TYPE_CHECKING, Any
+from typing import Literal, cast, overload
 
 import grpc
 
-if TYPE_CHECKING:
-    from hatchet_sdk.loader import ClientConfig
+from hatchet_sdk.config import ClientConfig
 
 
-def new_conn(config: "ClientConfig", aio=False):
+@overload
+def new_conn(config: ClientConfig, aio: Literal[False]) -> grpc.Channel: ...
 
+
+@overload
+def new_conn(config: ClientConfig, aio: Literal[True]) -> grpc.aio.Channel: ...
+
+
+def new_conn(config: ClientConfig, aio: bool) -> grpc.Channel | grpc.aio.Channel:
     credentials: grpc.ChannelCredentials | None = None
 
     # load channel credentials
-    if config.tls_config.tls_strategy == "tls":
-        root: Any | None = None
+    if config.tls_config.strategy == "tls":
+        root: bytes | None = None
 
-        if config.tls_config.ca_file:
-            root = open(config.tls_config.ca_file, "rb").read()
+        if config.tls_config.root_ca_file:
+            root = open(config.tls_config.root_ca_file, "rb").read()
 
         credentials = grpc.ssl_channel_credentials(root_certificates=root)
-    elif config.tls_config.tls_strategy == "mtls":
-        root = open(config.tls_config.ca_file, "rb").read()
+    elif config.tls_config.strategy == "mtls":
+        assert config.tls_config.root_ca_file
+        assert config.tls_config.key_file
+        assert config.tls_config.cert_file
+
+        root = open(config.tls_config.root_ca_file, "rb").read()
         private_key = open(config.tls_config.key_file, "rb").read()
         certificate_chain = open(config.tls_config.cert_file, "rb").read()
 
@@ -32,7 +42,7 @@ def new_conn(config: "ClientConfig", aio=False):
 
     start = grpc if not aio else grpc.aio
 
-    channel_options = [
+    channel_options: list[tuple[str, str | int]] = [
         ("grpc.max_send_message_length", config.grpc_max_send_message_length),
         ("grpc.max_receive_message_length", config.grpc_max_recv_message_length),
         ("grpc.keepalive_time_ms", 10 * 1000),
@@ -46,7 +56,7 @@ def new_conn(config: "ClientConfig", aio=False):
     # When steps execute via os.fork, we see `TSI_DATA_CORRUPTED` errors.
     os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "False"
 
-    if config.tls_config.tls_strategy == "none":
+    if config.tls_config.strategy == "none":
         conn = start.insecure_channel(
             target=config.host_port,
             options=channel_options,
@@ -61,4 +71,8 @@ def new_conn(config: "ClientConfig", aio=False):
             credentials=credentials,
             options=channel_options,
         )
-    return conn
+
+    return cast(
+        grpc.Channel | grpc.aio.Channel,
+        conn,
+    )

hatchet_sdk/context/__init__.py

@@ -1 +0,0 @@
-from .context import Context

hatchet_sdk/context/context.py

@@ -2,43 +2,34 @@ import inspect
 import json
 import traceback
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import Any, Generic, Type, TypeVar, cast, overload
-from warnings import warn
+from typing import Any, cast
 
-from pydantic import BaseModel, StrictStr
+from pydantic import BaseModel
 
-from hatchet_sdk.clients.events import EventClient
-from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
-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
-from hatchet_sdk.contracts.workflows_pb2 import (
-    BulkTriggerWorkflowRequest,
-    TriggerWorkflowRequest,
-)
-from hatchet_sdk.utils.types import WorkflowValidator
-from hatchet_sdk.utils.typing import is_basemodel_subclass
-from hatchet_sdk.workflow_run import WorkflowRunRef
-
-from ..clients.admin import (
+from hatchet_sdk.clients.admin import (
     AdminClient,
     ChildTriggerWorkflowOptions,
     ChildWorkflowRunDict,
     TriggerWorkflowOptions,
     WorkflowRunDict,
 )
-from ..clients.dispatcher.dispatcher import (  # type: ignore[attr-defined]
+from hatchet_sdk.clients.dispatcher.dispatcher import (  # type: ignore[attr-defined]
     Action,
     DispatcherClient,
 )
-from ..logger import logger
+from hatchet_sdk.clients.events import EventClient
+from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
+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
+from hatchet_sdk.logger import logger
+from hatchet_sdk.utils.typing import JSONSerializableMapping, WorkflowValidator
+from hatchet_sdk.workflow_run import WorkflowRunRef
 
 DEFAULT_WORKFLOW_POLLING_INTERVAL = 5  # Seconds
 
-T = TypeVar("T", bound=BaseModel)
-
 
 def get_caller_file_path() -> str:
     caller_frame = inspect.stack()[2]
@@ -46,125 +37,15 @@ def get_caller_file_path() -> str:
     return caller_frame.filename
 
 
-class BaseContext:
-
-    action: Action
-    spawn_index: int
-
-    def _prepare_workflow_options(
-        self,
-        key: str | None = None,
-        options: ChildTriggerWorkflowOptions | None = None,
-        worker_id: str | None = None,
-    ) -> TriggerWorkflowOptions:
-        workflow_run_id = self.action.workflow_run_id
-        step_run_id = self.action.step_run_id
-
-        desired_worker_id = None
-        if options is not None and "sticky" in options and options["sticky"] == True:
-            desired_worker_id = worker_id
-
-        meta = None
-        if options is not None and "additional_metadata" in options:
-            meta = options["additional_metadata"]
-
-        ## TODO: Pydantic here to simplify this
-        trigger_options: TriggerWorkflowOptions = {
-            "parent_id": workflow_run_id,
-            "parent_step_run_id": step_run_id,
-            "child_key": key,
-            "child_index": self.spawn_index,
-            "additional_metadata": meta,
-            "desired_worker_id": desired_worker_id,
-        }
-
-        self.spawn_index += 1
-        return trigger_options
-
-
-class ContextAioImpl(BaseContext):
-    def __init__(
-        self,
-        action: Action,
-        dispatcher_client: DispatcherClient,
-        admin_client: AdminClient,
-        event_client: EventClient,
-        rest_client: RestApi,
-        workflow_listener: PooledWorkflowRunListener,
-        workflow_run_event_listener: RunEventListenerClient,
-        worker: WorkerContext,
-        namespace: str = "",
-    ):
-        self.action = action
-        self.dispatcher_client = dispatcher_client
-        self.admin_client = admin_client
-        self.event_client = event_client
-        self.rest_client = rest_client
-        self.workflow_listener = workflow_listener
-        self.workflow_run_event_listener = workflow_run_event_listener
-        self.namespace = namespace
-        self.spawn_index = -1
-        self.worker = worker
-
-    @tenacity_retry
-    async def spawn_workflow(
-        self,
-        workflow_name: str,
-        input: dict[str, Any] = {},
-        key: str | None = None,
-        options: ChildTriggerWorkflowOptions | None = None,
-    ) -> WorkflowRunRef:
-        worker_id = self.worker.id()
-        # if (
-        #     options is not None
-        #     and "sticky" in options
-        #     and options["sticky"] == True
-        #     and not self.worker.has_workflow(workflow_name)
-        # ):
-        #     raise Exception(
-        #         f"cannot run with sticky: workflow {workflow_name} is not registered on the worker"
-        #     )
-
-        trigger_options = self._prepare_workflow_options(key, options, worker_id)
+class StepRunError(BaseModel):
+    step_id: str
+    step_run_action_name: str
+    error: str
 
-        return await self.admin_client.aio.run_workflow(
-            workflow_name, input, trigger_options
-        )
 
-    @tenacity_retry
-    async def spawn_workflows(
-        self, child_workflow_runs: list[ChildWorkflowRunDict]
-    ) -> list[WorkflowRunRef]:
-
-        if len(child_workflow_runs) == 0:
-            raise Exception("no child workflows to spawn")
-
-        worker_id = self.worker.id()
-
-        bulk_trigger_workflow_runs: list[WorkflowRunDict] = []
-        for child_workflow_run in child_workflow_runs:
-            workflow_name = child_workflow_run["workflow_name"]
-            input = child_workflow_run["input"]
-
-            key = child_workflow_run.get("key")
-            options = child_workflow_run.get("options", {})
-
-            trigger_options = self._prepare_workflow_options(key, options, worker_id)
-
-            bulk_trigger_workflow_runs.append(
-                WorkflowRunDict(
-                    workflow_name=workflow_name, input=input, options=trigger_options
-                )
-            )
-
-        return await self.admin_client.aio.run_workflows(bulk_trigger_workflow_runs)
-
-
-class Context(BaseContext):
+class Context:
     spawn_index = -1
 
-    worker: WorkerContext
-
     def __init__(
         self,
         action: Action,
@@ -172,7 +53,7 @@ class Context(BaseContext):
         admin_client: AdminClient,
         event_client: EventClient,
         rest_client: RestApi,
-        workflow_listener: PooledWorkflowRunListener,
+        workflow_listener: PooledWorkflowRunListener | None,
         workflow_run_event_listener: RunEventListenerClient,
         worker: WorkerContext,
         namespace: str = "",
@@ -181,38 +62,14 @@ class Context(BaseContext):
         self.worker = worker
         self.validator_registry = validator_registry
 
-        self.aio = ContextAioImpl(
-            action,
-            dispatcher_client,
-            admin_client,
-            event_client,
-            rest_client,
-            workflow_listener,
-            workflow_run_event_listener,
-            worker,
-            namespace,
-        )
-
-        # Check the type of action.action_payload before attempting to load it as JSON
-        if isinstance(action.action_payload, (str, bytes, bytearray)):
-            try:
-                self.data = cast(dict[str, Any], json.loads(action.action_payload))
-            except Exception as e:
-                logger.error(f"Error parsing action payload: {e}")
-                # Assign an empty dictionary if parsing fails
-                self.data: dict[str, Any] = {}  # type: ignore[no-redef]
-        else:
-            # Directly assign the payload to self.data if it's already a dict
-            self.data = (
-                action.action_payload if isinstance(action.action_payload, dict) else {}
-            )
+        self.data = action.action_payload
 
         self.action = action
 
         # FIXME: stepRunId is a legacy field, we should remove it
         self.stepRunId = action.step_run_id
 
-        self.step_run_id = action.step_run_id
+        self.step_run_id: str = action.step_run_id
         self.exit_flag = False
         self.dispatcher_client = dispatcher_client
         self.admin_client = admin_client
@@ -227,14 +84,28 @@ class Context(BaseContext):
         self.logger_thread_pool = ThreadPoolExecutor(max_workers=1)
         self.stream_event_thread_pool = ThreadPoolExecutor(max_workers=1)
 
-        # store each key in the overrides field in a lookup table
-        # overrides_data is a dictionary of key-value pairs
-        self.overrides_data = self.data.get("overrides", {})
+        self.input = self.data.input
 
-        if action.get_group_key_run_id != "":
-            self.input = self.data
-        else:
-            self.input = self.data.get("input", {})
+    def _prepare_workflow_options(
+        self,
+        key: str | None = None,
+        options: ChildTriggerWorkflowOptions = ChildTriggerWorkflowOptions(),
+        worker_id: str | None = None,
+    ) -> TriggerWorkflowOptions:
+        workflow_run_id = self.action.workflow_run_id
+        step_run_id = self.action.step_run_id
+
+        trigger_options = TriggerWorkflowOptions(
+            parent_id=workflow_run_id,
+            parent_step_run_id=step_run_id,
+            child_key=key,
+            child_index=self.spawn_index,
+            additional_metadata=options.additional_metadata,
+            desired_worker_id=worker_id if options.sticky else None,
+        )
+
+        self.spawn_index += 1
+        return trigger_options
 
     def step_output(self, step: str) -> dict[str, Any] | BaseModel:
         workflow_validator = next(
@@ -243,7 +114,7 @@ class Context(BaseContext):
         )
 
         try:
-            parent_step_data = cast(dict[str, Any], self.data["parents"][step])
+            parent_step_data = cast(dict[str, Any], self.data.parents[step])
         except KeyError:
             raise ValueError(f"Step output for '{step}' not found")
 
@@ -252,20 +123,15 @@ class Context(BaseContext):
 
         return parent_step_data
 
+    @property
     def triggered_by_event(self) -> bool:
-        return cast(str, self.data.get("triggered_by", "")) == "event"
-
-    def workflow_input(self) -> dict[str, Any] | T:
-        if (r := self.validator_registry.get(self.action.action_id)) and (
-            i := r.workflow_input
-        ):
-            return cast(
-                T,
-                i.model_validate(self.input),
-            )
+        return self.data.triggered_by == "event"
 
+    @property
+    def workflow_input(self) -> JSONSerializableMapping:
         return self.input
 
+    @property
     def workflow_run_id(self) -> str:
         return self.action.workflow_run_id
 
@@ -278,15 +144,6 @@ class Context(BaseContext):
         return self.exit_flag
 
     def playground(self, name: str, default: str | None = None) -> str | None:
-        # if the key exists in the overrides_data field, return the value
-        if name in self.overrides_data:
-            warn(
-                "Use of `overrides_data` is deprecated.",
-                DeprecationWarning,
-                stacklevel=1,
-            )
-            return str(self.overrides_data[name])
-
         caller_file = get_caller_file_path()
 
         self.dispatcher_client.put_overrides_data(
@@ -308,7 +165,9 @@ class Context(BaseContext):
             # we don't want to raise an exception here, as it will kill the log thread
             return False, e
 
-    def log(self, line: Any, raise_on_error: bool = False) -> None:
+    def log(
+        self, line: str | JSONSerializableMapping, raise_on_error: bool = False
+    ) -> None:
         if self.stepRunId == "":
             return
 
@@ -361,23 +220,29 @@ class Context(BaseContext):
         except Exception as e:
             logger.error(f"Error refreshing timeout: {e}")
 
+    @property
     def retry_count(self) -> int:
         return self.action.retry_count
 
-    def additional_metadata(self) -> dict[str, Any] | None:
+    @property
+    def additional_metadata(self) -> JSONSerializableMapping | None:
         return self.action.additional_metadata
 
+    @property
     def child_index(self) -> int | None:
         return self.action.child_workflow_index
 
+    @property
     def child_key(self) -> str | None:
         return self.action.child_workflow_key
 
+    @property
     def parent_workflow_run_id(self) -> str | None:
         return self.action.parent_workflow_run_id
 
+    @property
     def step_run_errors(self) -> dict[str, str]:
-        errors = cast(dict[str, str], self.data.get("step_run_errors", {}))
+        errors = self.data.step_run_errors
 
         if not errors:
             logger.error(
@@ -386,18 +251,18 @@ class Context(BaseContext):
 
         return errors
 
-    def fetch_run_failures(self) -> list[dict[str, StrictStr]]:
+    def fetch_run_failures(self) -> list[StepRunError]:
         data = self.rest_client.workflow_run_get(self.action.workflow_run_id)
         other_job_runs = [
             run for run in (data.job_runs or []) if run.job_id != self.action.job_id
         ]
-        # TODO: Parse Step Runs using a Pydantic Model rather than a hand crafted dictionary
+
         return [
-            {
-                "step_id": step_run.step_id,
-                "step_run_action_name": step_run.step.action,
-                "error": step_run.error,
-            }
+            StepRunError(
+                step_id=step_run.step_id,
+                step_run_action_name=step_run.step.action,
+                error=step_run.error,
+            )
             for job_run in other_job_runs
             if job_run.step_runs
             for step_run in job_run.step_runs
@@ -405,20 +270,23 @@ class Context(BaseContext):
         ]
 
     @tenacity_retry
-    def spawn_workflow(
+    async def aio_spawn_workflow(
         self,
         workflow_name: str,
-        input: dict[str, Any] = {},
+        input: JSONSerializableMapping = {},
         key: str | None = None,
-        options: ChildTriggerWorkflowOptions | None = None,
+        options: ChildTriggerWorkflowOptions = ChildTriggerWorkflowOptions(),
     ) -> WorkflowRunRef:
         worker_id = self.worker.id()
+
         trigger_options = self._prepare_workflow_options(key, options, worker_id)
 
-        return self.admin_client.run_workflow(workflow_name, input, trigger_options)
+        return await self.admin_client.aio_run_workflow(
+            workflow_name, input, trigger_options
+        )
 
     @tenacity_retry
-    def spawn_workflows(
+    async def aio_spawn_workflows(
         self, child_workflow_runs: list[ChildWorkflowRunDict]
     ) -> list[WorkflowRunRef]:
 
@@ -427,20 +295,52 @@ class Context(BaseContext):
 
         worker_id = self.worker.id()
 
-        bulk_trigger_workflow_runs: list[WorkflowRunDict] = []
-        for child_workflow_run in child_workflow_runs:
-            workflow_name = child_workflow_run["workflow_name"]
-            input = child_workflow_run["input"]
+        bulk_trigger_workflow_runs = [
+            WorkflowRunDict(
+                workflow_name=child_workflow_run.workflow_name,
+                input=child_workflow_run.input,
+                options=self._prepare_workflow_options(
+                    child_workflow_run.key, child_workflow_run.options, worker_id
+                ),
+            )
+            for child_workflow_run in child_workflow_runs
+        ]
 
-            key = child_workflow_run.get("key")
-            options = child_workflow_run.get("options", {})
+        return await self.admin_client.aio_run_workflows(bulk_trigger_workflow_runs)
 
-            trigger_options = self._prepare_workflow_options(key, options, worker_id)
+    @tenacity_retry
+    def spawn_workflow(
+        self,
+        workflow_name: str,
+        input: JSONSerializableMapping = {},
+        key: str | None = None,
+        options: ChildTriggerWorkflowOptions = ChildTriggerWorkflowOptions(),
+    ) -> WorkflowRunRef:
+        worker_id = self.worker.id()
+
+        trigger_options = self._prepare_workflow_options(key, options, worker_id)
+
+        return self.admin_client.run_workflow(workflow_name, input, trigger_options)
+
+    @tenacity_retry
+    def spawn_workflows(
+        self, child_workflow_runs: list[ChildWorkflowRunDict]
+    ) -> list[WorkflowRunRef]:
+
+        if len(child_workflow_runs) == 0:
+            raise Exception("no child workflows to spawn")
+
+        worker_id = self.worker.id()
 
-            bulk_trigger_workflow_runs.append(
-                WorkflowRunDict(
-                    workflow_name=workflow_name, input=input, options=trigger_options
-                )
+        bulk_trigger_workflow_runs = [
+            WorkflowRunDict(
+                workflow_name=child_workflow_run.workflow_name,
+                input=child_workflow_run.input,
+                options=self._prepare_workflow_options(
+                    child_workflow_run.key, child_workflow_run.options, worker_id
+                ),
             )
+            for child_workflow_run in child_workflow_runs
+        ]
 
         return self.admin_client.run_workflows(bulk_trigger_workflow_runs)