hatchet-sdk 1.14.1__tar.gz → 1.14.2__tar.gz

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/PKG-INFO

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

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/clients/v1/api_client.py

@@ -1,8 +1,13 @@
+from collections.abc import Callable
 from typing import ParamSpec, TypeVar
 
+import tenacity
+
 from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.configuration import Configuration
+from hatchet_sdk.clients.rest.exceptions import ServiceException
 from hatchet_sdk.config import ClientConfig
+from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 ## Type variables to use with coroutines.
@@ -20,7 +25,7 @@ P = ParamSpec("P")
 
 
 def maybe_additional_metadata_to_kv(
-    additional_metadata: dict[str, str] | JSONSerializableMapping | None
+    additional_metadata: dict[str, str] | JSONSerializableMapping | None,
 ) -> list[str] | None:
     if not additional_metadata:
         return None
@@ -42,3 +47,24 @@ class BaseRestClient:
 
     def client(self) -> ApiClient:
         return ApiClient(self.api_config)
+
+
+def retry(func: Callable[P, R]) -> Callable[P, R]:
+    return tenacity.retry(
+        reraise=True,
+        wait=tenacity.wait_exponential_jitter(),
+        stop=tenacity.stop_after_attempt(5),
+        before_sleep=_alert_on_retry,
+        retry=tenacity.retry_if_exception(_should_retry),
+    )(func)
+
+
+def _alert_on_retry(retry_state: tenacity.RetryCallState) -> None:
+    logger.debug(
+        f"Retrying {retry_state.fn}: attempt "
+        f"{retry_state.attempt_number} ended with: {retry_state.outcome}",
+    )
+
+
+def _should_retry(ex: BaseException) -> bool:
+    return isinstance(ex, ServiceException)

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/cron.py

@@ -19,6 +19,7 @@ from hatchet_sdk.clients.rest.models.workflow_run_order_by_direction import (
 from hatchet_sdk.clients.v1.api_client import (
     BaseRestClient,
     maybe_additional_metadata_to_kv,
+    retry,
 )
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
@@ -201,6 +202,7 @@ class CronClient(BaseRestClient):
             cron_name=cron_name,
         )
 
+    @retry
     def list(
         self,
         offset: int | None = None,
@@ -241,6 +243,7 @@ class CronClient(BaseRestClient):
                 cron_name=cron_name,
             )
 
+    @retry
     def get(self, cron_id: str) -> CronWorkflows:
         """
         Retrieve a specific workflow cron trigger by ID.

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/filters.py

@@ -10,7 +10,7 @@ from hatchet_sdk.clients.rest.models.v1_filter_list import V1FilterList
 from hatchet_sdk.clients.rest.models.v1_update_filter_request import (
     V1UpdateFilterRequest,
 )
-from hatchet_sdk.clients.v1.api_client import BaseRestClient
+from hatchet_sdk.clients.v1.api_client import BaseRestClient, retry
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 
@@ -41,6 +41,7 @@ class FiltersClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.list, limit, offset, workflow_ids, scopes)
 
+    @retry
     def list(
         self,
         limit: int | None = None,
@@ -67,6 +68,7 @@ class FiltersClient(BaseRestClient):
                 scopes=scopes,
             )
 
+    @retry
     def get(
         self,
         filter_id: str,

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/logs.py

@@ -3,7 +3,7 @@ import asyncio
 from hatchet_sdk.clients.rest.api.log_api import LogApi
 from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.models.v1_log_line_list import V1LogLineList
-from hatchet_sdk.clients.v1.api_client import BaseRestClient
+from hatchet_sdk.clients.v1.api_client import BaseRestClient, retry
 
 
 class LogsClient(BaseRestClient):
@@ -14,6 +14,7 @@ class LogsClient(BaseRestClient):
     def _la(self, client: ApiClient) -> LogApi:
         return LogApi(client)
 
+    @retry
     def list(self, task_run_id: str) -> V1LogLineList:
         """
         List log lines for a given task run.

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/metrics.py

@@ -12,6 +12,7 @@ from hatchet_sdk.clients.rest.models.workflow_run_status import WorkflowRunStatu
 from hatchet_sdk.clients.v1.api_client import (
     BaseRestClient,
     maybe_additional_metadata_to_kv,
+    retry,
 )
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
@@ -27,6 +28,7 @@ class MetricsClient(BaseRestClient):
     def _ta(self, client: ApiClient) -> TenantApi:
         return TenantApi(client)
 
+    @retry
     def get_workflow_metrics(
         self,
         workflow_id: str,
@@ -66,6 +68,7 @@ class MetricsClient(BaseRestClient):
             self.get_workflow_metrics, workflow_id, status, group_key
         )
 
+    @retry
     def get_queue_metrics(
         self,
         workflow_ids: list[str] | None = None,
@@ -105,6 +108,7 @@ class MetricsClient(BaseRestClient):
             self.get_queue_metrics, workflow_ids, additional_metadata
         )
 
+    @retry
     def get_task_metrics(self) -> TenantStepRunQueueMetrics:
         """
         Retrieve queue metrics

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/runs.py

@@ -1,10 +1,15 @@
 import asyncio
+from collections.abc import AsyncIterator
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Literal, overload
+from warnings import warn
 
 from pydantic import BaseModel, model_validator
 
-from hatchet_sdk.clients.listeners.run_event_listener import RunEventListenerClient
+from hatchet_sdk.clients.listeners.run_event_listener import (
+    RunEventListenerClient,
+    StepRunEventType,
+)
 from hatchet_sdk.clients.listeners.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.clients.rest.api.task_api import TaskApi
 from hatchet_sdk.clients.rest.api.workflow_runs_api import WorkflowRunsApi
@@ -13,6 +18,7 @@ from hatchet_sdk.clients.rest.models.v1_cancel_task_request import V1CancelTaskR
 from hatchet_sdk.clients.rest.models.v1_replay_task_request import V1ReplayTaskRequest
 from hatchet_sdk.clients.rest.models.v1_task_filter import V1TaskFilter
 from hatchet_sdk.clients.rest.models.v1_task_status import V1TaskStatus
+from hatchet_sdk.clients.rest.models.v1_task_summary import V1TaskSummary
 from hatchet_sdk.clients.rest.models.v1_task_summary_list import V1TaskSummaryList
 from hatchet_sdk.clients.rest.models.v1_trigger_workflow_run_request import (
     V1TriggerWorkflowRunRequest,
@@ -21,8 +27,11 @@ from hatchet_sdk.clients.rest.models.v1_workflow_run_details import V1WorkflowRu
 from hatchet_sdk.clients.v1.api_client import (
     BaseRestClient,
     maybe_additional_metadata_to_kv,
+    retry,
 )
 from hatchet_sdk.config import ClientConfig
+from hatchet_sdk.utils.aio import gather_max_concurrency
+from hatchet_sdk.utils.datetimes import partition_date_range
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
 if TYPE_CHECKING:
@@ -110,6 +119,7 @@ class RunsClient(BaseRestClient):
     def _ta(self, client: ApiClient) -> TaskApi:
         return TaskApi(client)
 
+    @retry
     def get(self, workflow_run_id: str) -> V1WorkflowRunDetails:
         """
         Get workflow run details for a given workflow run ID.
@@ -129,6 +139,7 @@ class RunsClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.get, workflow_run_id)
 
+    @retry
     def get_status(self, workflow_run_id: str) -> V1TaskStatus:
         """
         Get workflow run status for a given workflow run ID.
@@ -148,6 +159,152 @@ class RunsClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.get_status, workflow_run_id)
 
+    @retry
+    def list_with_pagination(
+        self,
+        since: datetime | None = None,
+        only_tasks: bool = False,
+        offset: int | None = None,
+        limit: int | None = None,
+        statuses: list[V1TaskStatus] | None = None,
+        until: datetime | None = None,
+        additional_metadata: dict[str, str] | None = None,
+        workflow_ids: list[str] | None = None,
+        worker_id: str | None = None,
+        parent_task_external_id: str | None = None,
+        triggering_event_external_id: str | None = None,
+    ) -> list[V1TaskSummary]:
+        """
+        List task runs according to a set of filters, paginating through days
+
+        :param since: The start time for filtering task runs.
+        :param only_tasks: Whether to only list task runs.
+        :param offset: The offset for pagination.
+        :param limit: The maximum number of task runs to return.
+        :param statuses: The statuses to filter task runs by.
+        :param until: The end time for filtering task runs.
+        :param additional_metadata: Additional metadata to filter task runs by.
+        :param workflow_ids: The workflow IDs to filter task runs by.
+        :param worker_id: The worker ID to filter task runs by.
+        :param parent_task_external_id: The parent task external ID to filter task runs by.
+        :param triggering_event_external_id: The event id that triggered the task run.
+
+        :return: A list of task runs matching the specified filters.
+        """
+
+        date_ranges = partition_date_range(
+            since=since or datetime.now(tz=timezone.utc) - timedelta(days=1),
+            until=until or datetime.now(tz=timezone.utc),
+        )
+
+        with self.client() as client:
+            responses = [
+                self._wra(client).v1_workflow_run_list(
+                    tenant=self.client_config.tenant_id,
+                    since=s,
+                    until=u,
+                    only_tasks=only_tasks,
+                    offset=offset,
+                    limit=limit,
+                    statuses=statuses,
+                    additional_metadata=maybe_additional_metadata_to_kv(
+                        additional_metadata
+                    ),
+                    workflow_ids=workflow_ids,
+                    worker_id=worker_id,
+                    parent_task_external_id=parent_task_external_id,
+                    triggering_event_external_id=triggering_event_external_id,
+                )
+                for s, u in date_ranges
+            ]
+
+            ## Hack for uniqueness
+            run_id_to_run = {
+                run.metadata.id: run for record in responses for run in record.rows
+            }
+
+            return sorted(
+                run_id_to_run.values(),
+                key=lambda x: x.created_at,
+                reverse=True,
+            )
+
+    @retry
+    async def aio_list_with_pagination(
+        self,
+        since: datetime | None = None,
+        only_tasks: bool = False,
+        offset: int | None = None,
+        limit: int | None = None,
+        statuses: list[V1TaskStatus] | None = None,
+        until: datetime | None = None,
+        additional_metadata: dict[str, str] | None = None,
+        workflow_ids: list[str] | None = None,
+        worker_id: str | None = None,
+        parent_task_external_id: str | None = None,
+        triggering_event_external_id: str | None = None,
+    ) -> list[V1TaskSummary]:
+        """
+        List task runs according to a set of filters, paginating through days
+
+        :param since: The start time for filtering task runs.
+        :param only_tasks: Whether to only list task runs.
+        :param offset: The offset for pagination.
+        :param limit: The maximum number of task runs to return.
+        :param statuses: The statuses to filter task runs by.
+        :param until: The end time for filtering task runs.
+        :param additional_metadata: Additional metadata to filter task runs by.
+        :param workflow_ids: The workflow IDs to filter task runs by.
+        :param worker_id: The worker ID to filter task runs by.
+        :param parent_task_external_id: The parent task external ID to filter task runs by.
+        :param triggering_event_external_id: The event id that triggered the task run.
+
+        :return: A list of task runs matching the specified filters.
+        """
+
+        date_ranges = partition_date_range(
+            since=since or datetime.now(tz=timezone.utc) - timedelta(days=1),
+            until=until or datetime.now(tz=timezone.utc),
+        )
+
+        with self.client() as client:
+            coros = [
+                asyncio.to_thread(
+                    self._wra(client).v1_workflow_run_list,
+                    tenant=self.client_config.tenant_id,
+                    since=s,
+                    until=u,
+                    only_tasks=only_tasks,
+                    offset=offset,
+                    limit=limit,
+                    statuses=statuses,
+                    additional_metadata=maybe_additional_metadata_to_kv(
+                        additional_metadata
+                    ),
+                    workflow_ids=workflow_ids,
+                    worker_id=worker_id,
+                    parent_task_external_id=parent_task_external_id,
+                    triggering_event_external_id=triggering_event_external_id,
+                )
+                for s, u in date_ranges
+            ]
+
+            responses = await gather_max_concurrency(
+                *coros,
+                max_concurrency=3,
+            )
+
+            ## Hack for uniqueness
+            run_id_to_run = {
+                run.metadata.id: run for record in responses for run in record.rows
+            }
+
+            return sorted(
+                run_id_to_run.values(),
+                key=lambda x: x.created_at,
+                reverse=True,
+            )
+
     async def aio_list(
         self,
         since: datetime | None = None,
@@ -181,7 +338,7 @@ class RunsClient(BaseRestClient):
         """
         return await asyncio.to_thread(
             self.list,
-            since=since or datetime.now(tz=timezone.utc) - timedelta(days=1),
+            since=since,
             only_tasks=only_tasks,
             offset=offset,
             limit=limit,
@@ -194,6 +351,7 @@ class RunsClient(BaseRestClient):
             triggering_event_external_id=triggering_event_external_id,
         )
 
+    @retry
     def list(
         self,
         since: datetime | None = None,
@@ -225,10 +383,22 @@ class RunsClient(BaseRestClient):
 
         :return: A list of task runs matching the specified filters.
         """
+
+        since = since or datetime.now(tz=timezone.utc) - timedelta(days=1)
+        until = until or datetime.now(tz=timezone.utc)
+
+        if (until - since).days >= 7:
+            warn(
+                "Listing runs with a date range longer than 7 days may result in performance issues. "
+                "Consider using `list_with_pagination` or `aio_list_with_pagination` instead.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
         with self.client() as client:
             return self._wra(client).v1_workflow_run_list(
                 tenant=self.client_config.tenant_id,
-                since=since or datetime.now(tz=timezone.utc) - timedelta(days=1),
+                since=since,
                 only_tasks=only_tasks,
                 offset=offset,
                 limit=limit,
@@ -376,6 +546,7 @@ class RunsClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.bulk_cancel, opts)
 
+    @retry
     def get_result(self, run_id: str) -> JSONSerializableMapping:
         """
         Get the result of a workflow run by its external ID.
@@ -413,3 +584,13 @@ class RunsClient(BaseRestClient):
             workflow_run_listener=self.workflow_run_listener,
             runs_client=self,
         )
+
+    async def subscribe_to_stream(
+        self,
+        workflow_run_id: str,
+    ) -> AsyncIterator[str]:
+        ref = self.get_run_ref(workflow_run_id=workflow_run_id)
+
+        async for chunk in ref.stream():
+            if chunk.type == StepRunEventType.STEP_RUN_EVENT_TYPE_STREAM:
+                yield chunk.payload

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/scheduled.py

@@ -21,6 +21,7 @@ from hatchet_sdk.clients.rest.models.workflow_run_order_by_direction import (
 from hatchet_sdk.clients.v1.api_client import (
     BaseRestClient,
     maybe_additional_metadata_to_kv,
+    retry,
 )
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 
@@ -153,6 +154,7 @@ class ScheduledClient(BaseRestClient):
             statuses=statuses,
         )
 
+    @retry
     def list(
         self,
         offset: int | None = None,
@@ -193,6 +195,7 @@ class ScheduledClient(BaseRestClient):
                 statuses=statuses,
             )
 
+    @retry
     def get(self, scheduled_id: str) -> ScheduledWorkflows:
         """
         Retrieves a specific scheduled workflow by scheduled run trigger ID.

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/tenant.py

@@ -3,7 +3,7 @@ import asyncio
 from hatchet_sdk.clients.rest.api.tenant_api import TenantApi
 from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.models.tenant import Tenant
-from hatchet_sdk.clients.v1.api_client import BaseRestClient
+from hatchet_sdk.clients.v1.api_client import BaseRestClient, retry
 
 
 class TenantClient(BaseRestClient):
@@ -14,6 +14,7 @@ class TenantClient(BaseRestClient):
     def _ta(self, client: ApiClient) -> TenantApi:
         return TenantApi(client)
 
+    @retry
     def get(self) -> Tenant:
         """
         Get the current tenant.

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/workers.py

@@ -5,7 +5,7 @@ from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.models.update_worker_request import UpdateWorkerRequest
 from hatchet_sdk.clients.rest.models.worker import Worker
 from hatchet_sdk.clients.rest.models.worker_list import WorkerList
-from hatchet_sdk.clients.v1.api_client import BaseRestClient
+from hatchet_sdk.clients.v1.api_client import BaseRestClient, retry
 
 
 class WorkersClient(BaseRestClient):
@@ -16,6 +16,7 @@ class WorkersClient(BaseRestClient):
     def _wa(self, client: ApiClient) -> WorkerApi:
         return WorkerApi(client)
 
+    @retry
     def get(self, worker_id: str) -> Worker:
         """
         Get a worker by its ID.
@@ -35,6 +36,7 @@ class WorkersClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.get, worker_id)
 
+    @retry
     def list(
         self,
     ) -> WorkerList:

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/features/workflows.py

@@ -6,7 +6,7 @@ from hatchet_sdk.clients.rest.api_client import ApiClient
 from hatchet_sdk.clients.rest.models.workflow import Workflow
 from hatchet_sdk.clients.rest.models.workflow_list import WorkflowList
 from hatchet_sdk.clients.rest.models.workflow_version import WorkflowVersion
-from hatchet_sdk.clients.v1.api_client import BaseRestClient
+from hatchet_sdk.clients.v1.api_client import BaseRestClient, retry
 
 
 class WorkflowsClient(BaseRestClient):
@@ -31,6 +31,7 @@ class WorkflowsClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.get, workflow_id)
 
+    @retry
     def get(self, workflow_id: str) -> Workflow:
         """
         Get a workflow by its ID.
@@ -41,6 +42,7 @@ class WorkflowsClient(BaseRestClient):
         with self.client() as client:
             return self._wa(client).workflow_get(workflow_id)
 
+    @retry
     def list(
         self,
         workflow_name: str | None = None,
@@ -81,6 +83,7 @@ class WorkflowsClient(BaseRestClient):
         """
         return await asyncio.to_thread(self.list, workflow_name, limit, offset)
 
+    @retry
     def get_version(
         self, workflow_id: str, version: str | None = None
     ) -> WorkflowVersion:

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/runnables/workflow.py

@@ -1,6 +1,6 @@
 import asyncio
 from collections.abc import Callable
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_type_hints
 
@@ -310,9 +310,9 @@ class BaseWorkflow(Generic[TWorkflowInput]):
 
         :returns: A list of `V1TaskSummary` objects representing the runs of the workflow.
         """
-        response = self.client.runs.list(
+        return self.client.runs.list_with_pagination(
             workflow_ids=[self.id],
-            since=since or datetime.now(tz=timezone.utc) - timedelta(days=1),
+            since=since,
             only_tasks=only_tasks,
             offset=offset,
             limit=limit,
@@ -324,8 +324,6 @@ class BaseWorkflow(Generic[TWorkflowInput]):
             triggering_event_external_id=triggering_event_external_id,
         )
 
-        return response.rows
-
     async def aio_list_runs(
         self,
         since: datetime | None = None,
@@ -355,9 +353,9 @@ class BaseWorkflow(Generic[TWorkflowInput]):
 
         :returns: A list of `V1TaskSummary` objects representing the runs of the workflow.
         """
-        return await asyncio.to_thread(
-            self.list_runs,
-            since=since or datetime.now(tz=timezone.utc) - timedelta(days=1),
+        return await self.client.runs.aio_list_with_pagination(
+            workflow_ids=[self.id],
+            since=since,
             only_tasks=only_tasks,
             offset=offset,
             limit=limit,

hatchet_sdk-1.14.2/hatchet_sdk/utils/aio.py

@@ -0,0 +1,21 @@
+import asyncio
+from collections.abc import Coroutine
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+async def gather_max_concurrency(
+    *tasks: Coroutine[None, None, T],
+    max_concurrency: int,
+) -> list[T]:
+    sem = asyncio.Semaphore(max_concurrency)
+
+    async def task_wrapper(task: Coroutine[None, None, T]) -> T:
+        async with sem:
+            return await task
+
+    return await asyncio.gather(
+        *(task_wrapper(task) for task in tasks),
+        return_exceptions=False,
+    )

hatchet_sdk-1.14.2/hatchet_sdk/utils/datetimes.py

@@ -0,0 +1,35 @@
+from datetime import datetime, timedelta, timezone
+from typing import TypeVar
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def _to_utc(dt: datetime) -> datetime:
+    if not dt.tzinfo:
+        return dt.replace(tzinfo=timezone.utc)
+
+    return dt.astimezone(timezone.utc)
+
+
+def partition_date_range(
+    since: datetime, until: datetime
+) -> list[tuple[datetime, datetime]]:
+    since = _to_utc(since)
+    until = _to_utc(until)
+
+    ranges = []
+    current_start = since
+
+    while current_start < until:
+        next_day = (current_start + timedelta(days=1)).replace(
+            hour=0, minute=0, second=0, microsecond=0
+        )
+
+        current_end = min(next_day, until)
+
+        ranges.append((current_start, current_end))
+
+        current_start = next_day
+
+    return ranges

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/hatchet_sdk/worker/runner/runner.py

@@ -550,7 +550,7 @@ class Runner:
             return ""
 
         if isinstance(output, BaseModel):
-            output = output.model_dump()
+            output = output.model_dump(mode="json")
 
         if not isinstance(output, dict):
             raise IllegalTaskOutputError(

{hatchet_sdk-1.14.1 → hatchet_sdk-1.14.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "1.14.1"
+version = "1.14.2"
 description = ""
 authors = ["Alexander Belanger <alexander@hatchet.run>"]
 readme = "README.md"