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

hatchet_sdk/clients/run_event_listener.py

@@ -1,9 +1,12 @@
 import asyncio
 import json
-from typing import AsyncGenerator
+from enum import Enum
+from typing import Any, AsyncGenerator, Callable, Generator, cast
 
 import grpc
+from pydantic import BaseModel
 
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     RESOURCE_TYPE_STEP_RUN,
@@ -13,15 +16,13 @@ from hatchet_sdk.contracts.dispatcher_pb2 import (
     WorkflowEvent,
 )
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
-
-from ..loader import ClientConfig
-from ..metadata import get_metadata
+from hatchet_sdk.metadata import get_metadata
 
 DEFAULT_ACTION_LISTENER_RETRY_INTERVAL = 5  # seconds
 DEFAULT_ACTION_LISTENER_RETRY_COUNT = 5
 
 
-class StepRunEventType:
+class StepRunEventType(str, Enum):
     STEP_RUN_EVENT_TYPE_STARTED = "STEP_RUN_EVENT_TYPE_STARTED"
     STEP_RUN_EVENT_TYPE_COMPLETED = "STEP_RUN_EVENT_TYPE_COMPLETED"
     STEP_RUN_EVENT_TYPE_FAILED = "STEP_RUN_EVENT_TYPE_FAILED"
@@ -30,7 +31,7 @@ class StepRunEventType:
     STEP_RUN_EVENT_TYPE_STREAM = "STEP_RUN_EVENT_TYPE_STREAM"
 
 
-class WorkflowRunEventType:
+class WorkflowRunEventType(str, Enum):
     WORKFLOW_RUN_EVENT_TYPE_STARTED = "WORKFLOW_RUN_EVENT_TYPE_STARTED"
     WORKFLOW_RUN_EVENT_TYPE_COMPLETED = "WORKFLOW_RUN_EVENT_TYPE_COMPLETED"
     WORKFLOW_RUN_EVENT_TYPE_FAILED = "WORKFLOW_RUN_EVENT_TYPE_FAILED"
@@ -56,50 +57,36 @@ workflow_run_event_type_mapping = {
 }
 
 
-class StepRunEvent:
-    def __init__(self, type: StepRunEventType, payload: str):
-        self.type = type
-        self.payload = payload
-
-
-def new_listener(config: ClientConfig):
-    return RunEventListenerClient(config=config)
+class StepRunEvent(BaseModel):
+    type: StepRunEventType
+    payload: str
 
 
 class RunEventListener:
-
-    workflow_run_id: str = None
-    additional_meta_kv: tuple[str, str] = None
-
-    def __init__(self, client: DispatcherStub, token: str):
+    def __init__(
+        self,
+        client: DispatcherStub,
+        token: str,
+        workflow_run_id: str | None = None,
+        additional_meta_kv: tuple[str, str] | None = None,
+    ):
         self.client = client
         self.stop_signal = False
         self.token = token
 
-    @classmethod
-    def for_run_id(cls, workflow_run_id: str, client: DispatcherStub, token: str):
-        listener = RunEventListener(client, token)
-        listener.workflow_run_id = workflow_run_id
-        return listener
+        self.workflow_run_id = workflow_run_id
+        self.additional_meta_kv = additional_meta_kv
 
-    @classmethod
-    def for_additional_meta(
-        cls, key: str, value: str, client: DispatcherStub, token: str
-    ):
-        listener = RunEventListener(client, token)
-        listener.additional_meta_kv = (key, value)
-        return listener
-
-    def abort(self):
+    def abort(self) -> None:
         self.stop_signal = True
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncGenerator[StepRunEvent, None]:
         return self._generator()
 
-    async def __anext__(self):
+    async def __anext__(self) -> StepRunEvent:
         return await self._generator().__anext__()
 
-    def __iter__(self):
+    def __iter__(self) -> Generator[StepRunEvent, None, None]:
         try:
             loop = asyncio.get_event_loop()
         except RuntimeError as e:
@@ -145,15 +132,18 @@ class RunEventListener:
 
                         try:
                             if workflow_event.eventPayload:
+                                ## TODO: Should this be `dumps` instead?
                                 payload = json.loads(workflow_event.eventPayload)
-                        except Exception as e:
+                        except Exception:
                             payload = workflow_event.eventPayload
                             pass
 
+                        assert isinstance(payload, str)
+
                         yield StepRunEvent(type=eventType, payload=payload)
                     elif workflow_event.resourceType == RESOURCE_TYPE_WORKFLOW_RUN:
-                        if workflow_event.eventType in workflow_run_event_type_mapping:
-                            eventType = workflow_run_event_type_mapping[
+                        if workflow_event.eventType in step_run_event_type_mapping:
+                            workflowRunEventType = step_run_event_type_mapping[
                                 workflow_event.eventType
                             ]
                         else:
@@ -166,10 +156,12 @@ class RunEventListener:
                         try:
                             if workflow_event.eventPayload:
                                 payload = json.loads(workflow_event.eventPayload)
-                        except Exception as e:
+                        except Exception:
                             pass
 
-                        yield StepRunEvent(type=eventType, payload=payload)
+                        assert isinstance(payload, str)
+
+                        yield StepRunEvent(type=workflowRunEventType, payload=payload)
 
                     if workflow_event.hangup:
                         listener = None
@@ -194,7 +186,7 @@ class RunEventListener:
                     break
                 # Raise StopAsyncIteration to properly end the generator
 
-    async def retry_subscribe(self):
+    async def retry_subscribe(self) -> AsyncGenerator[WorkflowEvent, None]:
         retries = 0
 
         while retries < DEFAULT_ACTION_LISTENER_RETRY_COUNT:
@@ -203,19 +195,25 @@ class RunEventListener:
                     await asyncio.sleep(DEFAULT_ACTION_LISTENER_RETRY_INTERVAL)
 
                 if self.workflow_run_id is not None:
-                    return self.client.SubscribeToWorkflowEvents(
-                        SubscribeToWorkflowEventsRequest(
-                            workflowRunId=self.workflow_run_id,
+                    return cast(
+                        AsyncGenerator[WorkflowEvent, None],
+                        self.client.SubscribeToWorkflowEvents(
+                            SubscribeToWorkflowEventsRequest(
+                                workflowRunId=self.workflow_run_id,
+                            ),
+                            metadata=get_metadata(self.token),
                         ),
-                        metadata=get_metadata(self.token),
                     )
                 elif self.additional_meta_kv is not None:
-                    return self.client.SubscribeToWorkflowEvents(
-                        SubscribeToWorkflowEventsRequest(
-                            additionalMetaKey=self.additional_meta_kv[0],
-                            additionalMetaValue=self.additional_meta_kv[1],
+                    return cast(
+                        AsyncGenerator[WorkflowEvent, None],
+                        self.client.SubscribeToWorkflowEvents(
+                            SubscribeToWorkflowEventsRequest(
+                                additionalMetaKey=self.additional_meta_kv[0],
+                                additionalMetaValue=self.additional_meta_kv[1],
+                            ),
+                            metadata=get_metadata(self.token),
                         ),
-                        metadata=get_metadata(self.token),
                     )
                 else:
                     raise Exception("no listener method provided")
@@ -226,34 +224,42 @@ class RunEventListener:
                 else:
                     raise ValueError(f"gRPC error: {e}")
 
+        raise Exception("Failed to subscribe to workflow events")
+
 
 class RunEventListenerClient:
     def __init__(self, config: ClientConfig):
         self.token = config.token
         self.config = config
-        self.client: DispatcherStub = None
+        self.client: DispatcherStub | None = None
 
-    def stream_by_run_id(self, workflow_run_id: str):
+    def stream_by_run_id(self, workflow_run_id: str) -> RunEventListener:
         return self.stream(workflow_run_id)
 
-    def stream(self, workflow_run_id: str):
+    def stream(self, workflow_run_id: str) -> RunEventListener:
         if not isinstance(workflow_run_id, str) and hasattr(workflow_run_id, "__str__"):
             workflow_run_id = str(workflow_run_id)
 
         if not self.client:
             aio_conn = new_conn(self.config, True)
-            self.client = DispatcherStub(aio_conn)
+            self.client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
 
-        return RunEventListener.for_run_id(workflow_run_id, self.client, self.token)
+        return RunEventListener(
+            client=self.client, token=self.token, workflow_run_id=workflow_run_id
+        )
 
-    def stream_by_additional_metadata(self, key: str, value: str):
+    def stream_by_additional_metadata(self, key: str, value: str) -> RunEventListener:
         if not self.client:
             aio_conn = new_conn(self.config, True)
-            self.client = DispatcherStub(aio_conn)
+            self.client = DispatcherStub(aio_conn)  # type: ignore[no-untyped-call]
 
-        return RunEventListener.for_additional_meta(key, value, self.client, self.token)
+        return RunEventListener(
+            client=self.client, token=self.token, additional_meta_kv=(key, value)
+        )
 
-    async def on(self, workflow_run_id: str, handler: callable = None):
+    async def on(
+        self, workflow_run_id: str, handler: Callable[[StepRunEvent], Any] | None = None
+    ) -> None:
         async for event in self.stream(workflow_run_id):
             # call the handler if provided
             if handler:

hatchet_sdk/clients/workflow_listener.py

@@ -1,22 +1,22 @@
 import asyncio
 import json
 from collections.abc import AsyncIterator
-from typing import AsyncGenerator
+from typing import Any, cast
 
 import grpc
-from grpc._cython import cygrpc
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
-from hatchet_sdk.clients.event_ts import Event_ts, read_with_interrupt
+from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.connection import new_conn
 from hatchet_sdk.contracts.dispatcher_pb2 import (
     SubscribeToWorkflowRunsRequest,
     WorkflowRunEvent,
 )
 from hatchet_sdk.contracts.dispatcher_pb2_grpc import DispatcherStub
-
-from ..loader import ClientConfig
-from ..logger import logger
-from ..metadata import get_metadata
+from hatchet_sdk.logger import logger
+from hatchet_sdk.metadata import get_metadata
 
 DEFAULT_WORKFLOW_LISTENER_RETRY_INTERVAL = 3  # seconds
 DEFAULT_WORKFLOW_LISTENER_RETRY_COUNT = 5
@@ -31,10 +31,10 @@ class _Subscription:
         self.workflow_run_id = workflow_run_id
         self.queue: asyncio.Queue[WorkflowRunEvent | None] = asyncio.Queue()
 
-    async def __aiter__(self):
+    async def __aiter__(self) -> "_Subscription":
         return self
 
-    async def __anext__(self) -> WorkflowRunEvent:
+    async def __anext__(self) -> WorkflowRunEvent | None:
         return await self.queue.get()
 
     async def get(self) -> WorkflowRunEvent:
@@ -45,35 +45,14 @@ class _Subscription:
 
         return event
 
-    async def put(self, item: WorkflowRunEvent):
+    async def put(self, item: WorkflowRunEvent) -> None:
         await self.queue.put(item)
 
-    async def close(self):
+    async def close(self) -> None:
         await self.queue.put(None)
 
 
 class PooledWorkflowRunListener:
-    # list of all active subscriptions, mapping from a subscription id to a workflow run id
-    subscriptionsToWorkflows: dict[int, str] = {}
-
-    # list of workflow run ids mapped to an array of subscription ids
-    workflowsToSubscriptions: dict[str, list[int]] = {}
-
-    subscription_counter: int = 0
-    subscription_counter_lock: asyncio.Lock = asyncio.Lock()
-
-    requests: asyncio.Queue[SubscribeToWorkflowRunsRequest] = asyncio.Queue()
-
-    listener: AsyncGenerator[WorkflowRunEvent, None] = None
-    listener_task: asyncio.Task = None
-
-    curr_requester: int = 0
-
-    # events have keys of the format workflow_run_id + subscription_id
-    events: dict[int, _Subscription] = {}
-
-    interrupter: asyncio.Task = None
-
     def __init__(self, config: ClientConfig):
         try:
             asyncio.get_running_loop()
@@ -82,11 +61,37 @@ class PooledWorkflowRunListener:
             asyncio.set_event_loop(loop)
 
         conn = new_conn(config, True)
-        self.client = DispatcherStub(conn)
+        self.client = DispatcherStub(conn)  # type: ignore[no-untyped-call]
         self.token = config.token
         self.config = config
 
-    async def _interrupter(self):
+        # list of all active subscriptions, mapping from a subscription id to a workflow run id
+        self.subscriptions_to_workflows: dict[int, str] = {}
+
+        # list of workflow run ids mapped to an array of subscription ids
+        self.workflows_to_subscriptions: dict[str, list[int]] = {}
+
+        self.subscription_counter: int = 0
+        self.subscription_counter_lock: asyncio.Lock = asyncio.Lock()
+
+        self.requests: asyncio.Queue[SubscribeToWorkflowRunsRequest | int] = (
+            asyncio.Queue()
+        )
+
+        self.listener: (
+            grpc.aio.UnaryStreamCall[SubscribeToWorkflowRunsRequest, WorkflowRunEvent]
+            | None
+        ) = None
+        self.listener_task: asyncio.Task[None] | None = None
+
+        self.curr_requester: int = 0
+
+        # events have keys of the format workflow_run_id + subscription_id
+        self.events: dict[int, _Subscription] = {}
+
+        self.interrupter: asyncio.Task[None] | None = None
+
+    async def _interrupter(self) -> None:
         """
         _interrupter runs in a separate thread and interrupts the listener according to a configurable duration.
         """
@@ -95,7 +100,7 @@ class PooledWorkflowRunListener:
         if self.interrupt is not None:
             self.interrupt.set()
 
-    async def _init_producer(self):
+    async def _init_producer(self) -> None:
         try:
             if not self.listener:
                 while True:
@@ -111,7 +116,10 @@ class PooledWorkflowRunListener:
                         self.interrupter = asyncio.create_task(self._interrupter())
 
                         while True:
-                            self.interrupt = Event_ts()
+                            self.interrupt = ThreadSafeEvent()
+                            if self.listener is None:
+                                continue
+
                             t = asyncio.create_task(
                                 read_with_interrupt(self.listener, self.interrupt)
                             )
@@ -124,7 +132,8 @@ class PooledWorkflowRunListener:
                                 )
 
                                 t.cancel()
-                                self.listener.cancel()
+                                if self.listener:
+                                    self.listener.cancel()
                                 await asyncio.sleep(
                                     DEFAULT_WORKFLOW_LISTENER_RETRY_INTERVAL
                                 )
@@ -136,7 +145,7 @@ class PooledWorkflowRunListener:
                                 break
 
                             # get a list of subscriptions for this workflow
-                            subscriptions = self.workflowsToSubscriptions.get(
+                            subscriptions = self.workflows_to_subscriptions.get(
                                 workflow_event.workflowRunId, []
                             )
 
@@ -163,7 +172,7 @@ class PooledWorkflowRunListener:
         self.curr_requester = self.curr_requester + 1
 
         # replay all existing subscriptions
-        workflow_run_set = set(self.subscriptionsToWorkflows.values())
+        workflow_run_set = set(self.subscriptions_to_workflows.values())
 
         for workflow_run_id in workflow_run_set:
             yield SubscribeToWorkflowRunsRequest(
@@ -184,17 +193,16 @@ class PooledWorkflowRunListener:
             yield request
             self.requests.task_done()
 
-    def cleanup_subscription(self, subscription_id: int):
-        workflow_run_id = self.subscriptionsToWorkflows[subscription_id]
+    def cleanup_subscription(self, subscription_id: int) -> None:
+        workflow_run_id = self.subscriptions_to_workflows[subscription_id]
 
-        if workflow_run_id in self.workflowsToSubscriptions:
-            self.workflowsToSubscriptions[workflow_run_id].remove(subscription_id)
+        if workflow_run_id in self.workflows_to_subscriptions:
+            self.workflows_to_subscriptions[workflow_run_id].remove(subscription_id)
 
-        del self.subscriptionsToWorkflows[subscription_id]
+        del self.subscriptions_to_workflows[subscription_id]
         del self.events[subscription_id]
 
-    async def subscribe(self, workflow_run_id: str):
-        init_producer: asyncio.Task = None
+    async def subscribe(self, workflow_run_id: str) -> WorkflowRunEvent:
         try:
             # create a new subscription id, place a mutex on the counter
             await self.subscription_counter_lock.acquire()
@@ -202,12 +210,12 @@ class PooledWorkflowRunListener:
             subscription_id = self.subscription_counter
             self.subscription_counter_lock.release()
 
-            self.subscriptionsToWorkflows[subscription_id] = workflow_run_id
+            self.subscriptions_to_workflows[subscription_id] = workflow_run_id
 
-            if workflow_run_id not in self.workflowsToSubscriptions:
-                self.workflowsToSubscriptions[workflow_run_id] = [subscription_id]
+            if workflow_run_id not in self.workflows_to_subscriptions:
+                self.workflows_to_subscriptions[workflow_run_id] = [subscription_id]
             else:
-                self.workflowsToSubscriptions[workflow_run_id].append(subscription_id)
+                self.workflows_to_subscriptions[workflow_run_id].append(subscription_id)
 
             self.events[subscription_id] = _Subscription(
                 subscription_id, workflow_run_id
@@ -222,23 +230,17 @@ class PooledWorkflowRunListener:
             if not self.listener_task or self.listener_task.done():
                 self.listener_task = asyncio.create_task(self._init_producer())
 
-            event = await self.events[subscription_id].get()
-
-            return event
+            return await self.events[subscription_id].get()
         except asyncio.CancelledError:
             raise
         finally:
             self.cleanup_subscription(subscription_id)
 
-    async def result(self, workflow_run_id: str):
+    async def result(self, workflow_run_id: str) -> dict[str, Any]:
         from hatchet_sdk.clients.admin import DedupeViolationErr
 
         event = await self.subscribe(workflow_run_id)
-
-        errors = []
-
-        if event.results:
-            errors = [result.error for result in event.results if result.error]
+        errors = [result.error for result in event.results if result.error]
 
         if errors:
             if DEDUPE_MESSAGE in errors[0]:
@@ -246,15 +248,15 @@ class PooledWorkflowRunListener:
             else:
                 raise Exception(f"Workflow Errors: {errors}")
 
-        results = {
+        return {
             result.stepReadableId: json.loads(result.output)
             for result in event.results
             if result.output
         }
 
-        return results
-
-    async def _retry_subscribe(self):
+    async def _retry_subscribe(
+        self,
+    ) -> grpc.aio.UnaryStreamCall[SubscribeToWorkflowRunsRequest, WorkflowRunEvent]:
         retries = 0
 
         while retries < DEFAULT_WORKFLOW_LISTENER_RETRY_COUNT:
@@ -266,12 +268,19 @@ class PooledWorkflowRunListener:
                 if self.curr_requester != 0:
                     self.requests.put_nowait(self.curr_requester)
 
-                return self.client.SubscribeToWorkflowRuns(
-                    self._request(),
-                    metadata=get_metadata(self.token),
+                return cast(
+                    grpc.aio.UnaryStreamCall[
+                        SubscribeToWorkflowRunsRequest, WorkflowRunEvent
+                    ],
+                    self.client.SubscribeToWorkflowRuns(
+                        self._request(),
+                        metadata=get_metadata(self.token),
+                    ),
                 )
             except grpc.RpcError as e:
                 if e.code() == grpc.StatusCode.UNAVAILABLE:
                     retries = retries + 1
                 else:
                     raise ValueError(f"gRPC error: {e}")
+
+        raise ValueError("Failed to connect to workflow run listener")

hatchet_sdk/config.py

@@ -0,0 +1,117 @@
+import json
+from logging import Logger, getLogger
+
+from pydantic import Field, field_validator, model_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from hatchet_sdk.token import get_addresses_from_jwt, get_tenant_id_from_jwt
+
+
+def create_settings_config(env_prefix: str) -> SettingsConfigDict:
+    return SettingsConfigDict(
+        env_prefix=env_prefix,
+        env_file=(".env", ".env.hatchet", ".env.dev", ".env.local"),
+        extra="ignore",
+    )
+
+
+class ClientTLSConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_TLS_",
+    )
+
+    strategy: str = "tls"
+    cert_file: str | None = None
+    key_file: str | None = None
+    root_ca_file: str | None = None
+    server_name: str = ""
+
+
+class HealthcheckConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_WORKER_HEALTHCHECK_",
+    )
+
+    port: int = 8001
+    enabled: bool = False
+
+
+DEFAULT_HOST_PORT = "localhost:7070"
+
+
+class ClientConfig(BaseSettings):
+    model_config = create_settings_config(
+        env_prefix="HATCHET_CLIENT_",
+    )
+
+    token: str = ""
+    logger: Logger = getLogger()
+
+    tenant_id: str = ""
+    host_port: str = DEFAULT_HOST_PORT
+    server_url: str = "https://app.dev.hatchet-tools.com"
+    namespace: str = ""
+
+    tls_config: ClientTLSConfig = Field(default_factory=lambda: ClientTLSConfig())
+    healthcheck: HealthcheckConfig = Field(default_factory=lambda: HealthcheckConfig())
+
+    listener_v2_timeout: int | None = None
+    grpc_max_recv_message_length: int = Field(
+        default=4 * 1024 * 1024, description="4MB default"
+    )
+    grpc_max_send_message_length: int = Field(
+        default=4 * 1024 * 1024, description="4MB default"
+    )
+
+    worker_preset_labels: dict[str, str] = Field(default_factory=dict)
+    enable_force_kill_sync_threads: bool = False
+
+    @model_validator(mode="after")
+    def validate_token_and_tenant(self) -> "ClientConfig":
+        if not self.token:
+            raise ValueError("Token must be set")
+
+        if not self.tenant_id:
+            self.tenant_id = get_tenant_id_from_jwt(self.token)
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_addresses(self) -> "ClientConfig":
+        if self.host_port == DEFAULT_HOST_PORT:
+            server_url, grpc_broadcast_address = get_addresses_from_jwt(self.token)
+            self.host_port = grpc_broadcast_address
+            self.server_url = server_url
+        else:
+            self.server_url = self.host_port
+
+        if not self.tls_config.server_name:
+            self.tls_config.server_name = self.host_port.split(":")[0]
+
+        if not self.tls_config.server_name:
+            self.tls_config.server_name = "localhost"
+
+        return self
+
+    @field_validator("listener_v2_timeout")
+    @classmethod
+    def validate_listener_timeout(cls, value: int | None | str) -> int | None:
+        if value is None:
+            return None
+
+        if isinstance(value, int):
+            return value
+
+        return int(value)
+
+    @field_validator("namespace")
+    @classmethod
+    def validate_namespace(cls, namespace: str) -> str:
+        if not namespace:
+            return ""
+        if not namespace.endswith("_"):
+            namespace = f"{namespace}_"
+        return namespace.lower()
+
+    def __hash__(self) -> int:
+        return hash(json.dumps(self.model_dump(), default=str))