hatchet-sdk 1.15.3__tar.gz → 1.16.0__tar.gz

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hatchet-sdk
-Version: 1.15.3
+Version: 1.16.0
 Summary: 
 License: MIT
 Author: Alexander Belanger

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/clients/admin.py

@@ -23,6 +23,7 @@ from hatchet_sdk.metadata import get_metadata
 from hatchet_sdk.rate_limit import RateLimitDuration
 from hatchet_sdk.runnables.contextvars import (
     ctx_action_key,
+    ctx_additional_metadata,
     ctx_step_run_id,
     ctx_worker_id,
     ctx_workflow_run_id,
@@ -278,6 +279,7 @@ class AdminClient:
         step_run_id = ctx_step_run_id.get()
         worker_id = ctx_worker_id.get()
         action_key = ctx_action_key.get()
+        additional_metadata = ctx_additional_metadata.get() or {}
         spawn_index = workflow_spawn_indices[action_key] if action_key else 0
 
         ## Increment the spawn_index for the parent workflow
@@ -296,7 +298,7 @@ class AdminClient:
             parent_step_run_id=options.parent_step_run_id or step_run_id,
             child_key=options.child_key,
             child_index=child_index,
-            additional_metadata=options.additional_metadata,
+            additional_metadata={**additional_metadata, **options.additional_metadata},
             desired_worker_id=desired_worker_id,
             priority=options.priority,
             namespace=options.namespace,

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/clients/dispatcher/action_listener.py

@@ -140,15 +140,15 @@ class ActionListener:
                     # todo case on "recvmsg:Connection reset by peer" for updates?
                     if self.missed_heartbeats >= 3:
                         # we don't reraise the error here, as we don't want to stop the heartbeat thread
-                        logger.error(
-                            f"⛔️ failed heartbeat ({self.missed_heartbeats}): {e.details()}"
+                        logger.exception(
+                            f"⛔️ failed heartbeat ({self.missed_heartbeats})"
                         )
                     elif self.missed_heartbeats > 1:
                         logger.warning(
                             f"failed to send heartbeat ({self.missed_heartbeats}): {e.details()}"
                         )
                 else:
-                    logger.error(f"failed to send heartbeat: {e}")
+                    logger.exception("failed to send heartbeat")
 
                 if self.interrupt is not None:
                     self.interrupt.set()
@@ -195,7 +195,7 @@ class ActionListener:
 
                     if not t.done():
                         logger.warning(
-                            "Interrupted read_with_interrupt task of action listener"
+                            "interrupted read_with_interrupt task of action listener"
                         )
 
                         t.cancel()
@@ -206,7 +206,7 @@ class ActionListener:
                     result = t.result()
 
                     if isinstance(result, UnexpectedEOF):
-                        logger.debug("Handling EOF in Action Listener")
+                        logger.debug("handling EOF in Action Listener")
                         self.retries = self.retries + 1
                         break
 
@@ -222,8 +222,8 @@ class ActionListener:
                                 assigned_action.actionPayload
                             )
                         )
-                    except (ValueError, json.JSONDecodeError) as e:
-                        logger.error(f"Error decoding payload: {e}")
+                    except (ValueError, json.JSONDecodeError):
+                        logger.exception("error decoding payload")
 
                         action_payload = ActionPayload()
 
@@ -263,9 +263,9 @@ class ActionListener:
                 # Handle different types of errors
                 if e.code() == grpc.StatusCode.CANCELLED:
                     # Context cancelled, unsubscribe and close
-                    logger.debug("Context cancelled, closing listener")
+                    logger.debug("context cancelled, closing listener")
                 elif e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
-                    logger.info("Deadline exceeded, retrying subscription")
+                    logger.info("deadline exceeded, retrying subscription")
                 elif (
                     self.listen_strategy == "v2"
                     and e.code() == grpc.StatusCode.UNIMPLEMENTED
@@ -277,10 +277,10 @@ class ActionListener:
                 else:
                     # TODO retry
                     if e.code() == grpc.StatusCode.UNAVAILABLE:
-                        logger.error(f"action listener error: {e.details()}")
+                        logger.exception("action listener error")
                     else:
                         # Unknown error, report and break
-                        logger.error(f"action listener error: {e}")
+                        logger.exception("action listener error")
 
                     self.retries = self.retries + 1
 
@@ -346,8 +346,8 @@ class ActionListener:
 
         try:
             self.unregister()
-        except Exception as e:
-            logger.error(f"failed to unregister: {e}")
+        except Exception:
+            logger.exception("failed to unregister")
 
         if self.interrupt:  # type: ignore[truthy-bool]
             self.interrupt.set()

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/clients/event_ts.py

@@ -66,7 +66,7 @@ async def read_with_interrupt(
         result = cast(TResponse, await listener.read())
 
         if result is cygrpc.EOF:
-            logger.warning("Received EOF from engine")
+            logger.warning("received EOF from engine")
             return UnexpectedEOF()
 
         key = key_generator(result) if key_generator else ""

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/clients/listeners/pooled_listener.py

@@ -103,7 +103,7 @@ class PooledListener(Generic[R, T, L], ABC):
                     try:
                         self.listener = await self._retry_subscribe()
 
-                        logger.debug("Listener connected.")
+                        logger.debug("listener connected.")
 
                         # spawn an interrupter task
                         if self.interrupter is not None and not self.interrupter.done():
@@ -125,7 +125,7 @@ class PooledListener(Generic[R, T, L], ABC):
 
                             if not t.done():
                                 logger.warning(
-                                    "Interrupted read_with_interrupt task of listener"
+                                    "interrupted read_with_interrupt task of listener"
                                 )
 
                                 t.cancel()
@@ -138,7 +138,7 @@ class PooledListener(Generic[R, T, L], ABC):
 
                             if isinstance(event, UnexpectedEOF):
                                 logger.debug(
-                                    f"Handling EOF in Pooled Listener {self.__class__.__name__}"
+                                    f"handling EOF in Pooled Listener {self.__class__.__name__}"
                                 )
                                 break
 
@@ -153,7 +153,7 @@ class PooledListener(Generic[R, T, L], ABC):
                         continue
 
         except Exception as e:
-            logger.error(f"Error in listener: {e}")
+            logger.exception("error in listener")
 
             self.listener = None
 

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/clients/rest/tenacity_utils.py

@@ -23,7 +23,7 @@ def tenacity_retry(func: Callable[P, R]) -> Callable[P, R]:
 def tenacity_alert_retry(retry_state: tenacity.RetryCallState) -> None:
     """Called between tenacity retries."""
     logger.debug(
-        f"Retrying {retry_state.fn}: attempt "
+        f"retrying {retry_state.fn}: attempt "
         f"{retry_state.attempt_number} ended with: {retry_state.outcome}",
     )
 

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/clients/v1/api_client.py

@@ -5,7 +5,7 @@ import tenacity
 
 from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.configuration import Configuration
-from hatchet_sdk.clients.rest.exceptions import ServiceException
+from hatchet_sdk.clients.rest.exceptions import NotFoundException, ServiceException
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.typing import JSONSerializableMapping
@@ -61,10 +61,10 @@ def retry(func: Callable[P, R]) -> Callable[P, R]:
 
 def _alert_on_retry(retry_state: tenacity.RetryCallState) -> None:
     logger.debug(
-        f"Retrying {retry_state.fn}: attempt "
+        f"retrying {retry_state.fn}: attempt "
         f"{retry_state.attempt_number} ended with: {retry_state.outcome}",
     )
 
 
 def _should_retry(ex: BaseException) -> bool:
-    return isinstance(ex, ServiceException)
+    return isinstance(ex, ServiceException | NotFoundException)

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/context/context.py

@@ -236,8 +236,8 @@ class Context:
                 step_run_id=self.step_run_id,
                 index=ix,
             )
-        except Exception as e:
-            logger.error(f"Error putting stream event: {e}")
+        except Exception:
+            logger.exception("error putting stream event")
 
     async def aio_put_stream(self, data: str | bytes) -> None:
         """
@@ -262,8 +262,8 @@ class Context:
             return self.dispatcher_client.refresh_timeout(
                 step_run_id=self.step_run_id, increment_by=increment_by
             )
-        except Exception as e:
-            logger.error(f"Error refreshing timeout: {e}")
+        except Exception:
+            logger.exception("error refreshing timeout")
 
     @property
     def retry_count(self) -> int:
@@ -285,7 +285,7 @@ class Context:
         return self.retry_count + 1
 
     @property
-    def additional_metadata(self) -> JSONSerializableMapping | None:
+    def additional_metadata(self) -> JSONSerializableMapping:
         """
         The additional metadata sent with the current task run.
 
@@ -350,7 +350,7 @@ class Context:
 
         if not errors:
             logger.error(
-                "No step run errors found. `context.task_run_errors` is intended to be run in an on-failure step, and will only work on engine versions more recent than v0.53.10"
+                "no step run errors found. `context.task_run_errors` is intended to be run in an on-failure step, and will only work on engine versions more recent than v0.53.10"
             )
 
         return errors

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/features/runs.py

@@ -119,6 +119,26 @@ class RunsClient(BaseRestClient):
     def _ta(self, client: ApiClient) -> TaskApi:
         return TaskApi(client)
 
+    @retry
+    def get_task_run(self, task_run_id: str) -> V1TaskSummary:
+        """
+        Get task run details for a given task run ID.
+
+        :param task_run_id: The ID of the task run to retrieve details for.
+        :return: Task run details for the specified task run ID.
+        """
+        with self.client() as client:
+            return self._ta(client).v1_task_get(task_run_id)
+
+    async def aio_get_task_run(self, task_run_id: str) -> V1TaskSummary:
+        """
+        Get task run details for a given task run ID.
+
+        :param task_run_id: The ID of the task run to retrieve details for.
+        :return: Task run details for the specified task run ID.
+        """
+        return await asyncio.to_thread(self.get_task_run, task_run_id)
+
     @retry
     def get(self, workflow_run_id: str) -> V1WorkflowRunDetails:
         """
@@ -148,7 +168,7 @@ class RunsClient(BaseRestClient):
         :return: The task status
         """
         with self.client() as client:
-            return self._wra(client).v1_workflow_run_get_status(str(workflow_run_id))
+            return self._wra(client).v1_workflow_run_get_status(workflow_run_id)
 
     async def aio_get_status(self, workflow_run_id: str) -> V1TaskStatus:
         """

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/opentelemetry/instrumentor.py

@@ -64,7 +64,7 @@ OTEL_TRACEPARENT_KEY = "traceparent"
 
 def create_traceparent() -> str | None:
     logger.warning(
-        "As of SDK version 1.11.0, you no longer need to call `create_traceparent` manually. The traceparent will be automatically created by the instrumentor and injected into the metadata of actions and events when appropriate. This method will be removed in a future version.",
+        "as of SDK version 1.11.0, you no longer need to call `create_traceparent` manually. The traceparent will be automatically created by the instrumentor and injected into the metadata of actions and events when appropriate. This method will be removed in a future version.",
     )
     return _create_traceparent()
 
@@ -91,7 +91,7 @@ def parse_carrier_from_metadata(
     metadata: JSONSerializableMapping | None,
 ) -> Context | None:
     logger.warning(
-        "As of SDK version 1.11.0, you no longer need to call `parse_carrier_from_metadata` manually. This method will be removed in a future version.",
+        "as of SDK version 1.11.0, you no longer need to call `parse_carrier_from_metadata` manually. This method will be removed in a future version.",
     )
 
     return _parse_carrier_from_metadata(metadata)
@@ -133,7 +133,7 @@ def inject_traceparent_into_metadata(
     metadata: dict[str, str], traceparent: str | None = None
 ) -> dict[str, str]:
     logger.warning(
-        "As of SDK version 1.11.0, you no longer need to call `inject_traceparent_into_metadata` manually. The traceparent will automatically be injected by the instrumentor. This method will be removed in a future version.",
+        "as of SDK version 1.11.0, you no longer need to call `inject_traceparent_into_metadata` manually. The traceparent will automatically be injected by the instrumentor. This method will be removed in a future version.",
     )
 
     return _inject_traceparent_into_metadata(metadata, traceparent)

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/runnables/contextvars.py

@@ -4,6 +4,7 @@ from collections import Counter
 from contextvars import ContextVar
 
 from hatchet_sdk.runnables.action import ActionKey
+from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 ctx_workflow_run_id: ContextVar[str | None] = ContextVar(
     "ctx_workflow_run_id", default=None
@@ -13,6 +14,9 @@ ctx_action_key: ContextVar[ActionKey | None] = ContextVar(
 )
 ctx_step_run_id: ContextVar[str | None] = ContextVar("ctx_step_run_id", default=None)
 ctx_worker_id: ContextVar[str | None] = ContextVar("ctx_worker_id", default=None)
+ctx_additional_metadata: ContextVar[JSONSerializableMapping | None] = ContextVar(
+    "ctx_additional_metadata", default=None
+)
 
 workflow_spawn_indices = Counter[ActionKey]()
 spawn_index_lock = asyncio.Lock()

{hatchet_sdk-1.15.3 → hatchet_sdk-1.16.0}/hatchet_sdk/runnables/task.py

@@ -11,6 +11,7 @@ from hatchet_sdk.conditions import (
     flatten_conditions,
 )
 from hatchet_sdk.context.context import Context, DurableContext
+from hatchet_sdk.context.worker_context import WorkerContext
 from hatchet_sdk.contracts.v1.shared.condition_pb2 import TaskConditions
 from hatchet_sdk.contracts.v1.workflows_pb2 import (
     CreateTaskOpts,
@@ -19,6 +20,7 @@ from hatchet_sdk.contracts.v1.workflows_pb2 import (
 )
 from hatchet_sdk.runnables.types import (
     ConcurrencyExpression,
+    EmptyModel,
     R,
     StepType,
     TWorkflowInput,
@@ -30,9 +32,11 @@ from hatchet_sdk.utils.timedelta_to_expression import Duration, timedelta_to_exp
 from hatchet_sdk.utils.typing import (
     AwaitableLike,
     CoroutineLike,
+    JSONSerializableMapping,
     TaskIOValidator,
     is_basemodel_subclass,
 )
+from hatchet_sdk.worker.runner.utils.capture_logs import AsyncLogSender
 
 if TYPE_CHECKING:
     from hatchet_sdk.runnables.workflow import Workflow
@@ -186,3 +190,132 @@ class Task(Generic[TWorkflowInput, R]):
             sleep_conditions=sleep_conditions,
             user_event_conditions=user_events,
         )
+
+    def _create_mock_context(
+        self,
+        input: TWorkflowInput | None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan_context: Any = None,
+    ) -> Context | DurableContext:
+        from hatchet_sdk.runnables.action import Action, ActionPayload, ActionType
+
+        additional_metadata = additional_metadata or {}
+        parent_outputs = parent_outputs or {}
+
+        if input is None:
+            input = cast(TWorkflowInput, EmptyModel())
+
+        action_payload = ActionPayload(input=input.model_dump(), parents=parent_outputs)
+
+        action = Action(
+            tenant_id=self.workflow.client.config.tenant_id,
+            worker_id="mock-worker-id",
+            workflow_run_id="mock-workflow-run-id",
+            get_group_key_run_id="mock-get-group-key-run-id",
+            job_id="mock-job-id",
+            job_name="mock-job-name",
+            job_run_id="mock-job-run-id",
+            step_id="mock-step-id",
+            step_run_id="mock-step-run-id",
+            action_id="mock:action",
+            action_payload=action_payload,
+            action_type=ActionType.START_STEP_RUN,
+            retry_count=retry_count,
+            additional_metadata=additional_metadata,
+            child_workflow_index=None,
+            child_workflow_key=None,
+            parent_workflow_run_id=None,
+            priority=1,
+            workflow_version_id="mock-workflow-version-id",
+            workflow_id="mock-workflow-id",
+        )
+
+        constructor = DurableContext if self.is_durable else Context
+
+        return constructor(
+            action=action,
+            dispatcher_client=self.workflow.client._client.dispatcher,
+            admin_client=self.workflow.client._client.admin,
+            event_client=self.workflow.client._client.event,
+            durable_event_listener=None,
+            worker=WorkerContext(
+                labels={}, client=self.workflow.client._client.dispatcher
+            ),
+            runs_client=self.workflow.client._client.runs,
+            lifespan_context=lifespan_context,
+            log_sender=AsyncLogSender(self.workflow.client._client.event),
+        )
+
+    def mock_run(
+        self,
+        input: TWorkflowInput | None = None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan: Any = None,
+    ) -> R:
+        """
+        Mimic the execution of a task. This method is intended to be used to unit test
+        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
+        tasks and `aio_mock_run` for async tasks.
+
+        :param input: The input to the task.
+        :param additional_metadata: Additional metadata to attach to the task.
+        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
+        :param retry_count: The number of times the task has been retried.
+        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
+
+        :return: The output of the task.
+        :raises TypeError: If the task is an async function and `mock_run` is called, or if the task is a sync function and `aio_mock_run` is called.
+        """
+
+        if self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not a sync function. Use `aio_mock_run` instead."
+            )
+
+        ctx = self._create_mock_context(
+            input, additional_metadata, parent_outputs, retry_count, lifespan
+        )
+
+        return self.call(ctx)
+
+    async def aio_mock_run(
+        self,
+        input: TWorkflowInput | None = None,
+        additional_metadata: JSONSerializableMapping | None = None,
+        parent_outputs: dict[str, JSONSerializableMapping] | None = None,
+        retry_count: int = 0,
+        lifespan: Any = None,
+    ) -> R:
+        """
+        Mimic the execution of a task. This method is intended to be used to unit test
+        tasks without needing to interact with the Hatchet engine. Use `mock_run` for sync
+        tasks and `aio_mock_run` for async tasks.
+
+        :param input: The input to the task.
+        :param additional_metadata: Additional metadata to attach to the task.
+        :param parent_outputs: Outputs from parent tasks, if any. This is useful for mimicking DAG functionality. For instance, if you have a task `step_2` that has a `parent` which is `step_1`, you can pass `parent_outputs={"step_1": {"result": "Hello, world!"}}` to `step_2.mock_run()` to be able to access `ctx.task_output(step_1)` in `step_2`.
+        :param retry_count: The number of times the task has been retried.
+        :param lifespan: The lifespan to be used in the task, which is useful if one was set on the worker. This will allow you to access `ctx.lifespan` inside of your task.
+
+        :return: The output of the task.
+        :raises TypeError: If the task is an async function and `mock_run` is called, or if the task is a sync function and `aio_mock_run` is called.
+        """
+
+        if not self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not an async function. Use `mock_run` instead."
+            )
+
+        ctx = self._create_mock_context(
+            input,
+            additional_metadata,
+            parent_outputs,
+            retry_count,
+            lifespan,
+        )
+
+        return await self.aio_call(ctx)