hatchet-sdk 1.2.6__py3-none-any.whl → 1.3.0__py3-none-any.whl

hatchet_sdk/__init__.py

@@ -3,8 +3,14 @@ from hatchet_sdk.clients.admin import (
     ScheduleTriggerWorkflowOptions,
     TriggerWorkflowOptions,
 )
-from hatchet_sdk.clients.durable_event_listener import RegisterDurableEventRequest
 from hatchet_sdk.clients.events import PushEventOptions
+from hatchet_sdk.clients.listeners.durable_event_listener import (
+    RegisterDurableEventRequest,
+)
+from hatchet_sdk.clients.listeners.run_event_listener import (
+    StepRunEventType,
+    WorkflowRunEventType,
+)
 from hatchet_sdk.clients.rest.models.accept_invite_request import AcceptInviteRequest
 
 # import models into sdk package
@@ -124,10 +130,6 @@ from hatchet_sdk.clients.rest.models.workflow_version_definition import (
     WorkflowVersionDefinition,
 )
 from hatchet_sdk.clients.rest.models.workflow_version_meta import WorkflowVersionMeta
-from hatchet_sdk.clients.run_event_listener import (
-    StepRunEventType,
-    WorkflowRunEventType,
-)
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.context.context import Context, DurableContext
 from hatchet_sdk.context.worker_context import WorkerContext

hatchet_sdk/client.py

@@ -1,8 +1,8 @@
 from hatchet_sdk.clients.admin import AdminClient
 from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
 from hatchet_sdk.clients.events import EventClient
-from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
-from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
+from hatchet_sdk.clients.listeners.run_event_listener import RunEventListenerClient
+from hatchet_sdk.clients.listeners.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.features.cron import CronClient
 from hatchet_sdk.features.logs import LogsClient
@@ -21,15 +21,15 @@ class Client:
         event_client: EventClient | None = None,
         admin_client: AdminClient | None = None,
         dispatcher_client: DispatcherClient | None = None,
-        workflow_listener: PooledWorkflowRunListener | None | None = None,
+        workflow_listener: PooledWorkflowRunListener | None = None,
         debug: bool = False,
     ):
         self.config = config
-        self.admin = admin_client or AdminClient(config)
         self.dispatcher = dispatcher_client or DispatcherClient(config)
         self.event = event_client or EventClient(config)
         self.listener = RunEventListenerClient(config)
-        self.workflow_listener = workflow_listener
+        self.workflow_listener = workflow_listener or PooledWorkflowRunListener(config)
+
         self.log_interceptor = config.logger
         self.debug = debug
 
@@ -37,7 +37,15 @@ class Client:
         self.logs = LogsClient(self.config)
         self.metrics = MetricsClient(self.config)
         self.rate_limits = RateLimitsClient(self.config)
-        self.runs = RunsClient(self.config)
+        self.runs = RunsClient(
+            config=self.config,
+            workflow_run_event_listener=self.listener,
+            workflow_run_listener=self.workflow_listener,
+        )
         self.scheduled = ScheduledClient(self.config)
         self.workers = WorkersClient(self.config)
         self.workflows = WorkflowsClient(self.config)
+
+        self.admin = admin_client or AdminClient(
+            config, self.workflow_listener, self.listener, self.runs
+        )

hatchet_sdk/clients/admin.py

@@ -7,6 +7,8 @@ import grpc
 from google.protobuf import timestamp_pb2
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
+from hatchet_sdk.clients.listeners.run_event_listener import RunEventListenerClient
+from hatchet_sdk.clients.listeners.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
 from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
@@ -14,6 +16,7 @@ from hatchet_sdk.contracts import workflows_pb2 as v0_workflow_protos
 from hatchet_sdk.contracts.v1 import workflows_pb2 as workflow_protos
 from hatchet_sdk.contracts.v1.workflows_pb2_grpc import AdminServiceStub
 from hatchet_sdk.contracts.workflows_pb2_grpc import WorkflowServiceStub
+from hatchet_sdk.features.runs import RunsClient
 from hatchet_sdk.metadata import get_metadata
 from hatchet_sdk.rate_limit import RateLimitDuration
 from hatchet_sdk.runnables.contextvars import (
@@ -63,14 +66,31 @@ class DedupeViolationErr(Exception):
 
 
 class AdminClient:
-    def __init__(self, config: ClientConfig):
-        conn = new_conn(config, False)
+    def __init__(
+        self,
+        config: ClientConfig,
+        workflow_run_listener: PooledWorkflowRunListener,
+        workflow_run_event_listener: RunEventListenerClient,
+        runs_client: RunsClient,
+    ):
         self.config = config
-        self.client = AdminServiceStub(conn)
-        self.v0_client = WorkflowServiceStub(conn)
+        self.runs_client = runs_client
         self.token = config.token
         self.namespace = config.namespace
 
+        self.workflow_run_listener = workflow_run_listener
+        self.workflow_run_event_listener = workflow_run_event_listener
+
+        self.client: AdminServiceStub | None = None
+        self.v0_client: WorkflowServiceStub | None = None
+
+    def _get_or_create_v0_client(self) -> WorkflowServiceStub:
+        if self.v0_client is None:
+            conn = new_conn(self.config, False)
+            self.v0_client = WorkflowServiceStub(conn)
+
+        return self.v0_client
+
     class TriggerWorkflowRequest(BaseModel):
         model_config = ConfigDict(extra="ignore")
 
@@ -199,6 +219,10 @@ class AdminClient:
     ) -> workflow_protos.CreateWorkflowVersionResponse:
         opts = self._prepare_put_workflow_request(name, workflow, overrides)
 
+        if self.client is None:
+            conn = new_conn(self.config, False)
+            self.client = AdminServiceStub(conn)
+
         return cast(
             workflow_protos.CreateWorkflowVersionResponse,
             self.client.PutWorkflow(
@@ -218,7 +242,9 @@ class AdminClient:
             duration, workflow_protos.RateLimitDuration
         )
 
-        self.v0_client.PutRateLimit(
+        client = self._get_or_create_v0_client()
+
+        client.PutRateLimit(
             v0_workflow_protos.PutRateLimitRequest(
                 key=key,
                 limit=limit,
@@ -245,9 +271,11 @@ class AdminClient:
                 name, schedules, input, options
             )
 
+            client = self._get_or_create_v0_client()
+
             return cast(
                 v0_workflow_protos.WorkflowVersion,
-                self.v0_client.ScheduleWorkflow(
+                client.ScheduleWorkflow(
                     request,
                     metadata=get_metadata(self.token),
                 ),
@@ -305,11 +333,12 @@ class AdminClient:
         options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
     ) -> WorkflowRunRef:
         request = self._create_workflow_run_request(workflow_name, input, options)
+        client = self._get_or_create_v0_client()
 
         try:
             resp = cast(
                 v0_workflow_protos.TriggerWorkflowResponse,
-                self.v0_client.TriggerWorkflow(
+                client.TriggerWorkflow(
                     request,
                     metadata=get_metadata(self.token),
                 ),
@@ -321,7 +350,9 @@ class AdminClient:
 
         return WorkflowRunRef(
             workflow_run_id=resp.workflow_run_id,
-            config=self.config,
+            workflow_run_event_listener=self.workflow_run_event_listener,
+            workflow_run_listener=self.workflow_run_listener,
+            runs_client=self.runs_client,
         )
 
     ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
@@ -332,13 +363,14 @@ class AdminClient:
         input: JSONSerializableMapping,
         options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
     ) -> WorkflowRunRef:
+        client = self._get_or_create_v0_client()
         async with spawn_index_lock:
             request = self._create_workflow_run_request(workflow_name, input, options)
 
         try:
             resp = cast(
                 v0_workflow_protos.TriggerWorkflowResponse,
-                self.v0_client.TriggerWorkflow(
+                client.TriggerWorkflow(
                     request,
                     metadata=get_metadata(self.token),
                 ),
@@ -350,8 +382,10 @@ class AdminClient:
             raise e
 
         return WorkflowRunRef(
+            runs_client=self.runs_client,
             workflow_run_id=resp.workflow_run_id,
-            config=self.config,
+            workflow_run_event_listener=self.workflow_run_event_listener,
+            workflow_run_listener=self.workflow_run_listener,
         )
 
     def chunk(self, xs: list[T], n: int) -> Generator[list[T], None, None]:
@@ -364,6 +398,7 @@ class AdminClient:
         self,
         workflows: list[WorkflowRunTriggerConfig],
     ) -> list[WorkflowRunRef]:
+        client = self._get_or_create_v0_client()
         bulk_workflows = [
             self._create_workflow_run_request(
                 workflow.workflow_name, workflow.input, workflow.options
@@ -380,7 +415,7 @@ class AdminClient:
 
             resp = cast(
                 v0_workflow_protos.BulkTriggerWorkflowResponse,
-                self.v0_client.BulkTriggerWorkflow(
+                client.BulkTriggerWorkflow(
                     bulk_request,
                     metadata=get_metadata(self.token),
                 ),
@@ -390,7 +425,9 @@ class AdminClient:
                 [
                     WorkflowRunRef(
                         workflow_run_id=workflow_run_id,
-                        config=self.config,
+                        workflow_run_event_listener=self.workflow_run_event_listener,
+                        workflow_run_listener=self.workflow_run_listener,
+                        runs_client=self.runs_client,
                     )
                     for workflow_run_id in resp.workflow_run_ids
                 ]
@@ -403,6 +440,7 @@ class AdminClient:
         self,
         workflows: list[WorkflowRunTriggerConfig],
     ) -> list[WorkflowRunRef]:
+        client = self._get_or_create_v0_client()
         chunks = self.chunk(workflows, MAX_BULK_WORKFLOW_RUN_BATCH_SIZE)
         refs: list[WorkflowRunRef] = []
 
@@ -421,7 +459,7 @@ class AdminClient:
 
             resp = cast(
                 v0_workflow_protos.BulkTriggerWorkflowResponse,
-                self.v0_client.BulkTriggerWorkflow(
+                client.BulkTriggerWorkflow(
                     bulk_request,
                     metadata=get_metadata(self.token),
                 ),
@@ -431,7 +469,9 @@ class AdminClient:
                 [
                     WorkflowRunRef(
                         workflow_run_id=workflow_run_id,
-                        config=self.config,
+                        workflow_run_event_listener=self.workflow_run_event_listener,
+                        workflow_run_listener=self.workflow_run_listener,
+                        runs_client=self.runs_client,
                     )
                     for workflow_run_id in resp.workflow_run_ids
                 ]
@@ -441,6 +481,8 @@ class AdminClient:
 
     def get_workflow_run(self, workflow_run_id: str) -> WorkflowRunRef:
         return WorkflowRunRef(
+            runs_client=self.runs_client,
             workflow_run_id=workflow_run_id,
-            config=self.config,
+            workflow_run_event_listener=self.workflow_run_event_listener,
+            workflow_run_listener=self.workflow_run_listener,
         )

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -12,7 +12,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_valida
 
 from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
 from hatchet_sdk.clients.events import proto_timestamp_now
-from hatchet_sdk.clients.run_event_listener import (
+from hatchet_sdk.clients.listeners.run_event_listener import (
     DEFAULT_ACTION_LISTENER_RETRY_INTERVAL,
 )
 from hatchet_sdk.config import ClientConfig
@@ -275,7 +275,7 @@ class ActionListener:
 
                         break
 
-                    assigned_action = t.result()
+                    assigned_action, _ = t.result()
 
                     if assigned_action is cygrpc.EOF:
                         self.retries = self.retries + 1

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -35,15 +35,20 @@ DEFAULT_REGISTER_TIMEOUT = 30
 
 class DispatcherClient:
     def __init__(self, config: ClientConfig):
-        conn = new_conn(config, False)
-        self.client = DispatcherStub(conn)
-
         self.token = config.token
         self.config = config
 
         ## IMPORTANT: This needs to be created lazily so we don't require
         ## an event loop to instantiate the client.
         self.aio_client: DispatcherStub | None = None
+        self.client: DispatcherStub | None = None
+
+    def _get_or_create_client(self) -> DispatcherStub:
+        if self.client is None:
+            conn = new_conn(self.config, False)
+            self.client = DispatcherStub(conn)
+
+        return self.client
 
     async def get_action_listener(
         self, req: GetActionListenerRequest
@@ -167,23 +172,29 @@ class DispatcherClient:
         )
 
     def put_overrides_data(self, data: OverridesData) -> ActionEventResponse:
+        client = self._get_or_create_client()
+
         return cast(
             ActionEventResponse,
-            self.client.PutOverridesData(
+            client.PutOverridesData(
                 data,
                 metadata=get_metadata(self.token),
             ),
         )
 
     def release_slot(self, step_run_id: str) -> None:
-        self.client.ReleaseSlot(
+        client = self._get_or_create_client()
+
+        client.ReleaseSlot(
             ReleaseSlotRequest(stepRunId=step_run_id),
             timeout=DEFAULT_REGISTER_TIMEOUT,
             metadata=get_metadata(self.token),
         )
 
     def refresh_timeout(self, step_run_id: str, increment_by: str) -> None:
-        self.client.RefreshTimeout(
+        client = self._get_or_create_client()
+
+        client.RefreshTimeout(
             RefreshTimeoutRequest(
                 stepRunId=step_run_id,
                 incrementTimeoutBy=increment_by,
@@ -203,7 +214,9 @@ class DispatcherClient:
             else:
                 worker_labels[key] = WorkerLabels(strValue=str(value))
 
-        self.client.UpsertWorkerLabels(
+        client = self._get_or_create_client()
+
+        client.UpsertWorkerLabels(
             UpsertWorkerLabelsRequest(workerId=worker_id, labels=worker_labels),
             timeout=DEFAULT_REGISTER_TIMEOUT,
             metadata=get_metadata(self.token),

hatchet_sdk/clients/event_ts.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import TypeVar, cast
+from typing import Callable, TypeVar, cast, overload
 
 import grpc.aio
 from grpc._cython import cygrpc  # type: ignore[attr-defined]
@@ -27,15 +27,35 @@ TRequest = TypeVar("TRequest")
 TResponse = TypeVar("TResponse")
 
 
+@overload
 async def read_with_interrupt(
-    listener: grpc.aio.UnaryStreamCall[TRequest, TResponse], interrupt: ThreadSafeEvent
-) -> TResponse:
+    listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
+    interrupt: ThreadSafeEvent,
+    key_generator: Callable[[TResponse], str],
+) -> tuple[TResponse, str]: ...
+
+
+@overload
+async def read_with_interrupt(
+    listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
+    interrupt: ThreadSafeEvent,
+    key_generator: None = None,
+) -> tuple[TResponse, None]: ...
+
+
+async def read_with_interrupt(
+    listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
+    interrupt: ThreadSafeEvent,
+    key_generator: Callable[[TResponse], str] | None = None,
+) -> tuple[TResponse, str | None]:
     try:
-        result = await listener.read()
+        result = cast(TResponse, await listener.read())
 
         if result is cygrpc.EOF:
             raise ValueError("Unexpected EOF")
 
-        return cast(TResponse, result)
+        key = key_generator(result) if key_generator else None
+
+        return result, key
     finally:
         interrupt.set()

hatchet_sdk/clients/listeners/durable_event_listener.py

@@ -0,0 +1,125 @@
+import json
+from collections.abc import AsyncIterator
+from typing import Any, Literal, cast
+
+import grpc
+import grpc.aio
+from pydantic import BaseModel, ConfigDict
+
+from hatchet_sdk.clients.listeners.pooled_listener import PooledListener
+from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
+from hatchet_sdk.connection import new_conn
+from hatchet_sdk.contracts.v1.dispatcher_pb2 import (
+    DurableEvent,
+    ListenForDurableEventRequest,
+)
+from hatchet_sdk.contracts.v1.dispatcher_pb2 import (
+    RegisterDurableEventRequest as RegisterDurableEventRequestProto,
+)
+from hatchet_sdk.contracts.v1.dispatcher_pb2_grpc import V1DispatcherStub
+from hatchet_sdk.contracts.v1.shared.condition_pb2 import DurableEventListenerConditions
+from hatchet_sdk.metadata import get_metadata
+from hatchet_sdk.waits import SleepCondition, UserEventCondition
+
+DEFAULT_DURABLE_EVENT_LISTENER_RETRY_INTERVAL = 3  # seconds
+DEFAULT_DURABLE_EVENT_LISTENER_RETRY_COUNT = 5
+DEFAULT_DURABLE_EVENT_LISTENER_INTERRUPT_INTERVAL = 1800  # 30 minutes
+
+
+class RegisterDurableEventRequest(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    task_id: str
+    signal_key: str
+    conditions: list[SleepCondition | UserEventCondition]
+
+    def to_proto(self) -> RegisterDurableEventRequestProto:
+        return RegisterDurableEventRequestProto(
+            task_id=self.task_id,
+            signal_key=self.signal_key,
+            conditions=DurableEventListenerConditions(
+                sleep_conditions=[
+                    c.to_proto()
+                    for c in self.conditions
+                    if isinstance(c, SleepCondition)
+                ],
+                user_event_conditions=[
+                    c.to_proto()
+                    for c in self.conditions
+                    if isinstance(c, UserEventCondition)
+                ],
+            ),
+        )
+
+
+class ParsedKey(BaseModel):
+    task_id: str
+    signal_key: str
+
+
+class DurableEventListener(
+    PooledListener[ListenForDurableEventRequest, DurableEvent, V1DispatcherStub]
+):
+    def _generate_key(self, task_id: str, signal_key: str) -> str:
+        return task_id + ":" + signal_key
+
+    def generate_key(self, response: DurableEvent) -> str:
+        return self._generate_key(
+            task_id=response.task_id,
+            signal_key=response.signal_key,
+        )
+
+    def parse_key(self, key: str) -> ParsedKey:
+        task_id, signal_key = key.split(":", maxsplit=1)
+
+        return ParsedKey(
+            task_id=task_id,
+            signal_key=signal_key,
+        )
+
+    async def create_subscription(
+        self,
+        request: AsyncIterator[ListenForDurableEventRequest],
+        metadata: tuple[tuple[str, str]],
+    ) -> grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent]:
+        if self.client is None:
+            conn = new_conn(self.config, True)
+            self.client = V1DispatcherStub(conn)
+
+        return cast(
+            grpc.aio.UnaryStreamCall[ListenForDurableEventRequest, DurableEvent],
+            self.client.ListenForDurableEvent(
+                request,  # type: ignore[arg-type]
+                metadata=metadata,
+            ),
+        )
+
+    def create_request_body(self, item: str) -> ListenForDurableEventRequest:
+        key = self.parse_key(item)
+        return ListenForDurableEventRequest(
+            task_id=key.task_id,
+            signal_key=key.signal_key,
+        )
+
+    @tenacity_retry
+    def register_durable_event(
+        self, request: RegisterDurableEventRequest
+    ) -> Literal[True]:
+        conn = new_conn(self.config, True)
+        client = V1DispatcherStub(conn)
+
+        client.RegisterDurableEvent(
+            request.to_proto(),
+            timeout=5,
+            metadata=get_metadata(self.token),
+        )
+
+        return True
+
+    @tenacity_retry
+    async def result(self, task_id: str, signal_key: str) -> dict[str, Any]:
+        key = self._generate_key(task_id, signal_key)
+
+        event = await self.subscribe(key)
+
+        return cast(dict[str, Any], json.loads(event.data.decode("utf-8")))