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

hatchet_sdk/clients/rest_client.py

@@ -2,7 +2,7 @@ import asyncio
 import atexit
 import datetime
 import threading
-from typing import Any, Coroutine, List
+from typing import Coroutine, TypeVar
 
 from pydantic import StrictInt
 
@@ -11,14 +11,13 @@ from hatchet_sdk.clients.rest.api.log_api import LogApi
 from hatchet_sdk.clients.rest.api.step_run_api import StepRunApi
 from hatchet_sdk.clients.rest.api.workflow_api import WorkflowApi
 from hatchet_sdk.clients.rest.api.workflow_run_api import WorkflowRunApi
-from hatchet_sdk.clients.rest.api.workflow_runs_api import WorkflowRunsApi
 from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.configuration import Configuration
-from hatchet_sdk.clients.rest.models import TriggerWorkflowRunRequest
 from hatchet_sdk.clients.rest.models.create_cron_workflow_trigger_request import (
     CreateCronWorkflowTriggerRequest,
 )
 from hatchet_sdk.clients.rest.models.cron_workflows import CronWorkflows
+from hatchet_sdk.clients.rest.models.cron_workflows_list import CronWorkflowsList
 from hatchet_sdk.clients.rest.models.cron_workflows_order_by_field import (
     CronWorkflowsOrderByField,
 )
@@ -27,6 +26,9 @@ from hatchet_sdk.clients.rest.models.event_order_by_direction import (
     EventOrderByDirection,
 )
 from hatchet_sdk.clients.rest.models.event_order_by_field import EventOrderByField
+from hatchet_sdk.clients.rest.models.event_update_cancel200_response import (
+    EventUpdateCancel200Response,
+)
 from hatchet_sdk.clients.rest.models.log_line_level import LogLineLevel
 from hatchet_sdk.clients.rest.models.log_line_list import LogLineList
 from hatchet_sdk.clients.rest.models.log_line_order_by_direction import (
@@ -44,16 +46,19 @@ from hatchet_sdk.clients.rest.models.schedule_workflow_run_request import (
     ScheduleWorkflowRunRequest,
 )
 from hatchet_sdk.clients.rest.models.scheduled_workflows import ScheduledWorkflows
+from hatchet_sdk.clients.rest.models.scheduled_workflows_list import (
+    ScheduledWorkflowsList,
+)
 from hatchet_sdk.clients.rest.models.scheduled_workflows_order_by_field import (
     ScheduledWorkflowsOrderByField,
 )
+from hatchet_sdk.clients.rest.models.trigger_workflow_run_request import (
+    TriggerWorkflowRunRequest,
+)
 from hatchet_sdk.clients.rest.models.workflow import Workflow
 from hatchet_sdk.clients.rest.models.workflow_kind import WorkflowKind
 from hatchet_sdk.clients.rest.models.workflow_list import WorkflowList
 from hatchet_sdk.clients.rest.models.workflow_run import WorkflowRun
-from hatchet_sdk.clients.rest.models.workflow_run_cancel200_response import (
-    WorkflowRunCancel200Response,
-)
 from hatchet_sdk.clients.rest.models.workflow_run_list import WorkflowRunList
 from hatchet_sdk.clients.rest.models.workflow_run_order_by_direction import (
     WorkflowRunOrderByDirection,
@@ -66,9 +71,21 @@ from hatchet_sdk.clients.rest.models.workflow_runs_cancel_request import (
     WorkflowRunsCancelRequest,
 )
 from hatchet_sdk.clients.rest.models.workflow_version import WorkflowVersion
+from hatchet_sdk.utils.typing import JSONSerializableMapping
+
+## Type variables to use with coroutines.
+## See https://stackoverflow.com/questions/73240620/the-right-way-to-type-hint-a-coroutine-function
+## Return type
+R = TypeVar("R")
 
+## Yield type
+Y = TypeVar("Y")
 
-class AsyncRestApi:
+## Send type
+S = TypeVar("S")
+
+
+class RestApi:
     def __init__(self, host: str, api_key: str, tenant_id: str):
         self.tenant_id = tenant_id
 
@@ -77,65 +94,85 @@ class AsyncRestApi:
             access_token=api_key,
         )
 
-        self._api_client = None
-        self._workflow_api = None
-        self._workflow_run_api = None
-        self._step_run_api = None
-        self._event_api = None
-        self._log_api = None
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(target=self._run_event_loop, daemon=True)
+        self._thread.start()
+
+        # Register the cleanup method to be called on exit
+        atexit.register(self._cleanup)
+
+        ## IMPORTANT: These clients need to be instantiated lazily because they rely on
+        ## an event loop to be running, which may not be the case when the `Hatchet` client is instantiated.
+        self._api_client: ApiClient | None = None
+        self._workflow_api: WorkflowApi | None = None
+        self._workflow_run_api: WorkflowRunApi | None = None
+        self._step_run_api: StepRunApi | None = None
+        self._event_api: EventApi | None = None
+        self._log_api: LogApi | None = None
 
     @property
-    def api_client(self):
+    def api_client(self) -> ApiClient:
+        ## IMPORTANT: This client needs to be instantiated lazily because it relies on an event
+        ## loop to be running, which may not be the case when the `Hatchet` client is instantiated.
         if self._api_client is None:
             self._api_client = ApiClient(configuration=self.config)
         return self._api_client
 
     @property
-    def workflow_api(self):
+    def workflow_api(self) -> WorkflowApi:
+        ## IMPORTANT: This client needs to be instantiated lazily because it relies on an event
+        ## loop to be running, which may not be the case when the `Hatchet` client is instantiated.
         if self._workflow_api is None:
             self._workflow_api = WorkflowApi(self.api_client)
         return self._workflow_api
 
     @property
-    def workflow_run_api(self):
+    def workflow_run_api(self) -> WorkflowRunApi:
+        ## IMPORTANT: This client needs to be instantiated lazily because it relies on an event
+        ## loop to be running, which may not be the case when the `Hatchet` client is instantiated.
         if self._workflow_run_api is None:
             self._workflow_run_api = WorkflowRunApi(self.api_client)
         return self._workflow_run_api
 
     @property
-    def step_run_api(self):
+    def step_run_api(self) -> StepRunApi:
+        ## IMPORTANT: This client needs to be instantiated lazily because it relies on an event
+        ## loop to be running, which may not be the case when the `Hatchet` client is instantiated.
         if self._step_run_api is None:
             self._step_run_api = StepRunApi(self.api_client)
         return self._step_run_api
 
     @property
-    def event_api(self):
+    def event_api(self) -> EventApi:
+        ## IMPORTANT: This client needs to be instantiated lazily because it relies on an event
+        ## loop to be running, which may not be the case when the `Hatchet` client is instantiated.
         if self._event_api is None:
             self._event_api = EventApi(self.api_client)
         return self._event_api
 
     @property
-    def log_api(self):
+    def log_api(self) -> LogApi:
+        ## IMPORTANT: This client needs to be instantiated lazily because it relies on an event
+        ## loop to be running, which may not be the case when the `Hatchet` client is instantiated.
         if self._log_api is None:
             self._log_api = LogApi(self.api_client)
         return self._log_api
 
-    async def close(self):
+    async def close(self) -> None:
         # Ensure the aiohttp client session is closed
-        if self._api_client is not None:
-            await self._api_client.close()
+        await self.api_client.close()  # type: ignore[no-untyped-call]
 
-    async def workflow_list(self) -> WorkflowList:
+    async def aio_list_workflows(self) -> WorkflowList:
         return await self.workflow_api.workflow_list(
             tenant=self.tenant_id,
         )
 
-    async def workflow_get(self, workflow_id: str) -> Workflow:
+    async def aio_get_workflow(self, workflow_id: str) -> Workflow:
         return await self.workflow_api.workflow_get(
             workflow=workflow_id,
         )
 
-    async def workflow_version_get(
+    async def aio_get_workflow_version(
         self, workflow_id: str, version: str | None = None
     ) -> WorkflowVersion:
         return await self.workflow_api.workflow_version_get(
@@ -143,7 +180,7 @@ class AsyncRestApi:
             version=version,
         )
 
-    async def workflow_run_list(
+    async def aio_list_workflow_runs(
         self,
         workflow_id: str | None = None,
         offset: int | None = None,
@@ -172,25 +209,25 @@ class AsyncRestApi:
             order_by_direction=order_by_direction,
         )
 
-    async def workflow_run_get(self, workflow_run_id: str) -> WorkflowRun:
+    async def aio_get_workflow_run(self, workflow_run_id: str) -> WorkflowRun:
         return await self.workflow_api.workflow_run_get(
             tenant=self.tenant_id,
             workflow_run=workflow_run_id,
         )
 
-    async def workflow_run_replay(
+    async def aio_replay_workflow_run(
         self, workflow_run_ids: list[str]
     ) -> ReplayWorkflowRunsResponse:
         return await self.workflow_run_api.workflow_run_update_replay(
             tenant=self.tenant_id,
             replay_workflow_runs_request=ReplayWorkflowRunsRequest(
-                workflow_run_ids=workflow_run_ids,
+                workflowRunIds=workflow_run_ids,
             ),
         )
 
-    async def workflow_run_cancel(
+    async def aio_cancel_workflow_run(
         self, workflow_run_id: str
-    ) -> WorkflowRunCancel200Response:
+    ) -> EventUpdateCancel200Response:
         return await self.workflow_run_api.workflow_run_cancel(
             tenant=self.tenant_id,
             workflow_runs_cancel_request=WorkflowRunsCancelRequest(
@@ -198,9 +235,9 @@ class AsyncRestApi:
             ),
         )
 
-    async def workflow_run_bulk_cancel(
+    async def aio_bulk_cancel_workflow_runs(
         self, workflow_run_ids: list[str]
-    ) -> WorkflowRunCancel200Response:
+    ) -> EventUpdateCancel200Response:
         return await self.workflow_run_api.workflow_run_cancel(
             tenant=self.tenant_id,
             workflow_runs_cancel_request=WorkflowRunsCancelRequest(
@@ -208,48 +245,48 @@ class AsyncRestApi:
             ),
         )
 
-    async def workflow_run_create(
+    async def aio_create_workflow_run(
         self,
         workflow_id: str,
-        input: dict[str, Any],
+        input: JSONSerializableMapping,
         version: str | None = None,
-        additional_metadata: list[str] | None = None,
+        additional_metadata: JSONSerializableMapping = {},
     ) -> WorkflowRun:
         return await self.workflow_run_api.workflow_run_create(
             workflow=workflow_id,
             version=version,
             trigger_workflow_run_request=TriggerWorkflowRunRequest(
-                input=input,
-                additional_metadata=additional_metadata,
+                input=dict(input),
+                additionalMetadata=dict(additional_metadata),
             ),
         )
 
-    async def cron_create(
+    async def aio_create_cron(
         self,
         workflow_name: str,
         cron_name: str,
         expression: str,
-        input: dict[str, Any],
-        additional_metadata: dict[str, str],
-    ):
+        input: JSONSerializableMapping,
+        additional_metadata: JSONSerializableMapping,
+    ) -> CronWorkflows:
         return await self.workflow_run_api.cron_workflow_trigger_create(
             tenant=self.tenant_id,
             workflow=workflow_name,
             create_cron_workflow_trigger_request=CreateCronWorkflowTriggerRequest(
                 cronName=cron_name,
                 cronExpression=expression,
-                input=input,
-                additional_metadata=additional_metadata,
+                input=dict(input),
+                additionalMetadata=dict(additional_metadata),
             ),
         )
 
-    async def cron_delete(self, cron_trigger_id: str):
-        return await self.workflow_api.workflow_cron_delete(
+    async def aio_delete_cron(self, cron_trigger_id: str) -> None:
+        await self.workflow_api.workflow_cron_delete(
             tenant=self.tenant_id,
             cron_workflow=cron_trigger_id,
         )
 
-    async def cron_list(
+    async def aio_list_crons(
         self,
         offset: StrictInt | None = None,
         limit: StrictInt | None = None,
@@ -257,7 +294,7 @@ class AsyncRestApi:
         additional_metadata: list[str] | None = None,
         order_by_field: CronWorkflowsOrderByField | None = None,
         order_by_direction: WorkflowRunOrderByDirection | None = None,
-    ):
+    ) -> CronWorkflowsList:
         return await self.workflow_api.cron_workflow_list(
             tenant=self.tenant_id,
             offset=offset,
@@ -268,36 +305,36 @@ class AsyncRestApi:
             order_by_direction=order_by_direction,
         )
 
-    async def cron_get(self, cron_trigger_id: str):
+    async def aio_get_cron(self, cron_trigger_id: str) -> CronWorkflows:
         return await self.workflow_api.workflow_cron_get(
             tenant=self.tenant_id,
             cron_workflow=cron_trigger_id,
         )
 
-    async def schedule_create(
+    async def aio_create_schedule(
         self,
         name: str,
         trigger_at: datetime.datetime,
-        input: dict[str, Any],
-        additional_metadata: dict[str, str],
-    ):
+        input: JSONSerializableMapping,
+        additional_metadata: JSONSerializableMapping,
+    ) -> ScheduledWorkflows:
         return await self.workflow_run_api.scheduled_workflow_run_create(
             tenant=self.tenant_id,
             workflow=name,
             schedule_workflow_run_request=ScheduleWorkflowRunRequest(
                 triggerAt=trigger_at,
-                input=input,
-                additional_metadata=additional_metadata,
+                input=dict(input),
+                additionalMetadata=dict(additional_metadata),
             ),
         )
 
-    async def schedule_delete(self, scheduled_trigger_id: str):
-        return await self.workflow_api.workflow_scheduled_delete(
+    async def aio_delete_schedule(self, scheduled_trigger_id: str) -> None:
+        await self.workflow_api.workflow_scheduled_delete(
             tenant=self.tenant_id,
             scheduled_workflow_run=scheduled_trigger_id,
         )
 
-    async def schedule_list(
+    async def aio_list_schedule(
         self,
         offset: StrictInt | None = None,
         limit: StrictInt | None = None,
@@ -307,7 +344,7 @@ class AsyncRestApi:
         parent_step_run_id: str | None = None,
         order_by_field: ScheduledWorkflowsOrderByField | None = None,
         order_by_direction: WorkflowRunOrderByDirection | None = None,
-    ):
+    ) -> ScheduledWorkflowsList:
         return await self.workflow_api.workflow_scheduled_list(
             tenant=self.tenant_id,
             offset=offset,
@@ -320,13 +357,13 @@ class AsyncRestApi:
             order_by_direction=order_by_direction,
         )
 
-    async def schedule_get(self, scheduled_trigger_id: str):
+    async def aio_get_schedule(self, scheduled_trigger_id: str) -> ScheduledWorkflows:
         return await self.workflow_api.workflow_scheduled_get(
             tenant=self.tenant_id,
             scheduled_workflow_run=scheduled_trigger_id,
         )
 
-    async def list_logs(
+    async def aio_list_logs(
         self,
         step_run_id: str,
         offset: int | None = None,
@@ -346,7 +383,7 @@ class AsyncRestApi:
             order_by_direction=order_by_direction,
         )
 
-    async def events_list(
+    async def aio_list_events(
         self,
         offset: int | None = None,
         limit: int | None = None,
@@ -371,44 +408,32 @@ class AsyncRestApi:
             additional_metadata=additional_metadata,
         )
 
-    async def events_replay(self, event_ids: list[str] | EventList) -> EventList:
+    async def aio_replay_events(self, event_ids: list[str] | EventList) -> EventList:
         if isinstance(event_ids, EventList):
-            event_ids = [r.metadata.id for r in event_ids.rows]
+            rows = event_ids.rows or []
+            event_ids = [r.metadata.id for r in rows]
 
-        return self.event_api.event_update_replay(
+        return await self.event_api.event_update_replay(
             tenant=self.tenant_id,
             replay_event_request=ReplayEventRequest(eventIds=event_ids),
         )
 
-
-class RestApi:
-    def __init__(self, host: str, api_key: str, tenant_id: str):
-        self._loop = asyncio.new_event_loop()
-        self._thread = threading.Thread(target=self._run_event_loop, daemon=True)
-        self._thread.start()
-
-        # Initialize AsyncRestApi inside the event loop to ensure an active loop
-        self.aio = AsyncRestApi(host, api_key, tenant_id)
-
-        # Register the cleanup method to be called on exit
-        atexit.register(self._cleanup)
-
-    def _cleanup(self):
+    def _cleanup(self) -> None:
         """
         Stop the running thread and clean up the event loop.
         """
-        self._run_coroutine(self.aio.close())
+        self._run_coroutine(self.close())
         self._loop.call_soon_threadsafe(self._loop.stop)
         self._thread.join()
 
-    def _run_event_loop(self):
+    def _run_event_loop(self) -> None:
         """
         Run the asyncio event loop in a separate thread.
         """
         asyncio.set_event_loop(self._loop)
         self._loop.run_forever()
 
-    def _run_coroutine(self, coro) -> Any:
+    def _run_coroutine(self, coro: Coroutine[Y, S, R]) -> R:
         """
         Execute a coroutine in the event loop and return the result.
         """
@@ -416,15 +441,15 @@ class RestApi:
         return future.result()
 
     def workflow_list(self) -> WorkflowList:
-        return self._run_coroutine(self.aio.workflow_list())
+        return self._run_coroutine(self.aio_list_workflows())
 
     def workflow_get(self, workflow_id: str) -> Workflow:
-        return self._run_coroutine(self.aio.workflow_get(workflow_id))
+        return self._run_coroutine(self.aio_get_workflow(workflow_id))
 
     def workflow_version_get(
         self, workflow_id: str, version: str | None = None
     ) -> WorkflowVersion:
-        return self._run_coroutine(self.aio.workflow_version_get(workflow_id, version))
+        return self._run_coroutine(self.aio_get_workflow_version(workflow_id, version))
 
     def workflow_run_list(
         self,
@@ -441,7 +466,7 @@ class RestApi:
         order_by_direction: WorkflowRunOrderByDirection | None = None,
     ) -> WorkflowRunList:
         return self._run_coroutine(
-            self.aio.workflow_run_list(
+            self.aio_list_workflow_runs(
                 workflow_id=workflow_id,
                 offset=offset,
                 limit=limit,
@@ -457,25 +482,25 @@ class RestApi:
         )
 
     def workflow_run_get(self, workflow_run_id: str) -> WorkflowRun:
-        return self._run_coroutine(self.aio.workflow_run_get(workflow_run_id))
+        return self._run_coroutine(self.aio_get_workflow_run(workflow_run_id))
 
-    def workflow_run_cancel(self, workflow_run_id: str) -> WorkflowRunCancel200Response:
-        return self._run_coroutine(self.aio.workflow_run_cancel(workflow_run_id))
+    def workflow_run_cancel(self, workflow_run_id: str) -> EventUpdateCancel200Response:
+        return self._run_coroutine(self.aio_cancel_workflow_run(workflow_run_id))
 
     def workflow_run_bulk_cancel(
         self, workflow_run_ids: list[str]
-    ) -> WorkflowRunCancel200Response:
-        return self._run_coroutine(self.aio.workflow_run_bulk_cancel(workflow_run_ids))
+    ) -> EventUpdateCancel200Response:
+        return self._run_coroutine(self.aio_bulk_cancel_workflow_runs(workflow_run_ids))
 
     def workflow_run_create(
         self,
         workflow_id: str,
-        input: dict[str, Any],
+        input: JSONSerializableMapping,
         version: str | None = None,
-        additional_metadata: list[str] | None = None,
+        additional_metadata: JSONSerializableMapping = {},
     ) -> WorkflowRun:
         return self._run_coroutine(
-            self.aio.workflow_run_create(
+            self.aio_create_workflow_run(
                 workflow_id, input, version, additional_metadata
             )
         )
@@ -485,17 +510,17 @@ class RestApi:
         workflow_name: str,
         cron_name: str,
         expression: str,
-        input: dict[str, Any],
-        additional_metadata: dict[str, str],
+        input: JSONSerializableMapping,
+        additional_metadata: JSONSerializableMapping,
     ) -> CronWorkflows:
         return self._run_coroutine(
-            self.aio.cron_create(
+            self.aio_create_cron(
                 workflow_name, cron_name, expression, input, additional_metadata
             )
         )
 
-    def cron_delete(self, cron_trigger_id: str):
-        return self._run_coroutine(self.aio.cron_delete(cron_trigger_id))
+    def cron_delete(self, cron_trigger_id: str) -> None:
+        self._run_coroutine(self.aio_delete_cron(cron_trigger_id))
 
     def cron_list(
         self,
@@ -505,9 +530,9 @@ class RestApi:
         additional_metadata: list[str] | None = None,
         order_by_field: CronWorkflowsOrderByField | None = None,
         order_by_direction: WorkflowRunOrderByDirection | None = None,
-    ):
+    ) -> CronWorkflowsList:
         return self._run_coroutine(
-            self.aio.cron_list(
+            self.aio_list_crons(
                 offset,
                 limit,
                 workflow_id,
@@ -517,24 +542,24 @@ class RestApi:
             )
         )
 
-    def cron_get(self, cron_trigger_id: str):
-        return self._run_coroutine(self.aio.cron_get(cron_trigger_id))
+    def cron_get(self, cron_trigger_id: str) -> CronWorkflows:
+        return self._run_coroutine(self.aio_get_cron(cron_trigger_id))
 
     def schedule_create(
         self,
         workflow_name: str,
         trigger_at: datetime.datetime,
-        input: dict[str, Any],
-        additional_metadata: dict[str, str],
-    ):
+        input: JSONSerializableMapping,
+        additional_metadata: JSONSerializableMapping,
+    ) -> ScheduledWorkflows:
         return self._run_coroutine(
-            self.aio.schedule_create(
+            self.aio_create_schedule(
                 workflow_name, trigger_at, input, additional_metadata
             )
         )
 
-    def schedule_delete(self, scheduled_trigger_id: str):
-        return self._run_coroutine(self.aio.schedule_delete(scheduled_trigger_id))
+    def schedule_delete(self, scheduled_trigger_id: str) -> None:
+        self._run_coroutine(self.aio_delete_schedule(scheduled_trigger_id))
 
     def schedule_list(
         self,
@@ -544,9 +569,9 @@ class RestApi:
         additional_metadata: list[str] | None = None,
         order_by_field: CronWorkflowsOrderByField | None = None,
         order_by_direction: WorkflowRunOrderByDirection | None = None,
-    ):
+    ) -> ScheduledWorkflowsList:
         return self._run_coroutine(
-            self.aio.schedule_list(
+            self.aio_list_schedule(
                 offset,
                 limit,
                 workflow_id,
@@ -556,8 +581,8 @@ class RestApi:
             )
         )
 
-    def schedule_get(self, scheduled_trigger_id: str):
-        return self._run_coroutine(self.aio.schedule_get(scheduled_trigger_id))
+    def schedule_get(self, scheduled_trigger_id: str) -> ScheduledWorkflows:
+        return self._run_coroutine(self.aio_get_schedule(scheduled_trigger_id))
 
     def list_logs(
         self,
@@ -570,7 +595,7 @@ class RestApi:
         order_by_direction: LogLineOrderByDirection | None = None,
     ) -> LogLineList:
         return self._run_coroutine(
-            self.aio.list_logs(
+            self.aio_list_logs(
                 step_run_id=step_run_id,
                 offset=offset,
                 limit=limit,
@@ -594,7 +619,7 @@ class RestApi:
         additional_metadata: list[str] | None = None,
     ) -> EventList:
         return self._run_coroutine(
-            self.aio.events_list(
+            self.aio_list_events(
                 offset=offset,
                 limit=limit,
                 keys=keys,
@@ -608,4 +633,4 @@ class RestApi:
         )
 
     def events_replay(self, event_ids: list[str] | EventList) -> EventList:
-        return self._run_coroutine(self.aio.events_replay(event_ids))
+        return self._run_coroutine(self.aio_replay_events(event_ids))