hatchet-sdk 0.47.1__py3-none-any.whl → 1.0.0a1__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
 
@@ -15,11 +15,11 @@ 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,
 )
@@ -28,6 +28,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 (
@@ -45,16 +48,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,
@@ -67,9 +73,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
 
@@ -78,73 +96,104 @@ 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._workflow_runs_api: WorkflowRunsApi | None = None
+        self._step_run_api: StepRunApi | None = None
+        self._event_api: EventApi | None = None
+        self._log_api: LogApi | None = None
         self._worker_api: WorkerApi | 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 worker_api(self):
+    def workflow_runs_api(self) -> WorkflowRunsApi:
+        ## 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_runs_api is None:
+            self._workflow_runs_api = WorkflowRunsApi(self.api_client)
+        return self._workflow_runs_api
+
+    @property
+    def worker_api(self) -> WorkerApi:
+        ## 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._worker_api is None:
             self._worker_api = WorkerApi(self.api_client)
 
         return self._worker_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(
@@ -152,7 +201,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,
@@ -181,25 +230,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(
@@ -207,9 +256,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(
@@ -217,48 +266,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,
@@ -266,7 +315,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,
@@ -277,36 +326,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,
@@ -316,7 +365,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,
@@ -329,13 +378,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,
@@ -355,7 +404,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,
@@ -380,44 +429,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.
         """
@@ -425,15 +462,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,
@@ -450,7 +487,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,
@@ -466,25 +503,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
             )
         )
@@ -494,17 +531,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,
@@ -514,9 +551,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,
@@ -526,24 +563,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,
@@ -553,9 +590,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,
@@ -565,8 +602,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,
@@ -579,7 +616,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,
@@ -603,7 +640,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,
@@ -617,4 +654,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))