hatchet-sdk 0.47.1__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

hatchet_sdk/labels.py

@@ -1,10 +1,8 @@
-from typing import TypedDict
+from pydantic import BaseModel
 
 
-class DesiredWorkerLabel(TypedDict, total=False):
+class DesiredWorkerLabel(BaseModel):
     value: str | int
-    required: bool | None = None
+    required: bool = False
     weight: int | None = None
-    comparator: int | None = (
-        None  # _ClassVar[WorkerLabelComparator] TODO figure out type
-    )
+    comparator: int | None = None

hatchet_sdk/metadata.py

@@ -1,2 +1,2 @@
-def get_metadata(token: str):
+def get_metadata(token: str) -> list[tuple[str, str]]:
     return [("authorization", "bearer " + token)]

hatchet_sdk/opentelemetry/instrumentor.py

@@ -1,6 +1,8 @@
 from importlib.metadata import version
 from typing import Any, Callable, Collection, Coroutine
 
+from hatchet_sdk.utils.typing import JSONSerializableMapping
+
 try:
     from opentelemetry.context import Context
     from opentelemetry.instrumentation.instrumentor import (  # type: ignore[attr-defined]
@@ -28,7 +30,7 @@ import hatchet_sdk
 from hatchet_sdk.clients.admin import (
     AdminClient,
     TriggerWorkflowOptions,
-    WorkflowRunDict,
+    WorkflowRunTriggerConfig,
 )
 from hatchet_sdk.clients.dispatcher.action_listener import Action
 from hatchet_sdk.clients.events import (
@@ -66,7 +68,9 @@ def create_traceparent() -> str | None:
     return carrier.get("traceparent")
 
 
-def parse_carrier_from_metadata(metadata: dict[str, str] | None) -> Context | None:
+def parse_carrier_from_metadata(
+    metadata: JSONSerializableMapping | None,
+) -> Context | None:
     """
     Parses OpenTelemetry trace context from a metadata dictionary.
 
@@ -202,7 +206,7 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
 
         wrap_function_wrapper(
             hatchet_sdk,
-            "clients.admin.AdminClientAioImpl.run_workflow",
+            "clients.admin.AdminClient.aio_run_workflow",
             self._wrap_async_run_workflow,
         )
 
@@ -214,7 +218,7 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
 
         wrap_function_wrapper(
             hatchet_sdk,
-            "clients.admin.AdminClientAioImpl.run_workflows",
+            "clients.admin.AdminClient.aio_run_workflows",
             self._wrap_async_run_workflows,
         )
 
@@ -348,14 +352,12 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
     def _wrap_run_workflows(
         self,
         wrapped: Callable[
-            [list[WorkflowRunDict], TriggerWorkflowOptions | None], list[WorkflowRunRef]
+            [list[WorkflowRunTriggerConfig]],
+            list[WorkflowRunRef],
         ],
         instance: AdminClient,
-        args: tuple[
-            list[WorkflowRunDict],
-            TriggerWorkflowOptions | None,
-        ],
-        kwargs: dict[str, list[WorkflowRunDict] | TriggerWorkflowOptions | None],
+        args: tuple[list[WorkflowRunTriggerConfig],],
+        kwargs: dict[str, list[WorkflowRunTriggerConfig]],
     ) -> list[WorkflowRunRef]:
         with self._tracer.start_as_current_span(
             "hatchet.run_workflows",
@@ -366,15 +368,12 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
     async def _wrap_async_run_workflows(
         self,
         wrapped: Callable[
-            [list[WorkflowRunDict], TriggerWorkflowOptions | None],
+            [list[WorkflowRunTriggerConfig]],
             Coroutine[None, None, list[WorkflowRunRef]],
         ],
         instance: AdminClient,
-        args: tuple[
-            list[WorkflowRunDict],
-            TriggerWorkflowOptions | None,
-        ],
-        kwargs: dict[str, list[WorkflowRunDict] | TriggerWorkflowOptions | None],
+        args: tuple[list[WorkflowRunTriggerConfig],],
+        kwargs: dict[str, list[WorkflowRunTriggerConfig]],
     ) -> list[WorkflowRunRef]:
         with self._tracer.start_as_current_span(
             "hatchet.run_workflows",
@@ -391,6 +390,6 @@ class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
         unwrap(hatchet_sdk, "clients.events.EventClient.push")
         unwrap(hatchet_sdk, "clients.events.EventClient.bulk_push")
         unwrap(hatchet_sdk, "clients.admin.AdminClient.run_workflow")
-        unwrap(hatchet_sdk, "clients.admin.AdminClientAioImpl.run_workflow")
+        unwrap(hatchet_sdk, "clients.admin.AdminClient.aio_run_workflow")
         unwrap(hatchet_sdk, "clients.admin.AdminClient.run_workflows")
-        unwrap(hatchet_sdk, "clients.admin.AdminClientAioImpl.run_workflows")
+        unwrap(hatchet_sdk, "clients.admin.AdminClient.aio_run_workflows")

hatchet_sdk/rate_limit.py

@@ -1,9 +1,9 @@
-from dataclasses import dataclass
-from typing import Union
+from enum import Enum
 
-from celpy import CELEvalError, Environment
+from celpy import CELEvalError, Environment  # type: ignore
+from pydantic import BaseModel, model_validator
 
-from hatchet_sdk.contracts.workflows_pb2 import CreateStepRateLimit
+from hatchet_sdk.contracts.v1.workflows_pb2 import CreateTaskRateLimit
 
 
 def validate_cel_expression(expr: str) -> bool:
@@ -15,7 +15,7 @@ def validate_cel_expression(expr: str) -> bool:
         return False
 
 
-class RateLimitDuration:
+class RateLimitDuration(str, Enum):
     SECOND = "SECOND"
     MINUTE = "MINUTE"
     HOUR = "HOUR"
@@ -25,8 +25,7 @@ class RateLimitDuration:
     YEAR = "YEAR"
 
 
-@dataclass
-class RateLimit:
+class RateLimit(BaseModel):
     """
     Represents a rate limit configuration for a step in a workflow.
 
@@ -40,7 +39,6 @@ class RateLimit:
         units (int or str, default=1): The number of units or a CEL expression for dynamic unit calculation.
         limit (int or str, optional): The rate limit value or a CEL expression for dynamic limit calculation.
         duration (str, default=RateLimitDuration.MINUTE): The window duration of the rate limit.
-        key (str, optional): Deprecated. Use static_key instead.
 
     Usage:
         1. Static rate limit:
@@ -61,62 +59,49 @@ class RateLimit:
 
     Raises:
         ValueError: If invalid combinations of attributes are provided or if CEL expressions are invalid.
-        DeprecationWarning: If the deprecated 'key' attribute is used.
     """
 
-    key: Union[str, None] = None
-    static_key: Union[str, None] = None
-    dynamic_key: Union[str, None] = None
-    units: Union[int, str] = 1
-    limit: Union[int, str, None] = None
+    static_key: str | None = None
+    dynamic_key: str | None = None
+    units: str | int = 1
+    limit: int | str | None = None
     duration: RateLimitDuration = RateLimitDuration.MINUTE
 
-    _req: CreateStepRateLimit = None
+    @model_validator(mode="after")
+    def validate_rate_limit(self) -> "RateLimit":
+        if self.dynamic_key and self.static_key:
+            raise ValueError("Cannot have both static key and dynamic key set")
 
-    def __post_init__(self):
-        # juggle the key and key_expr fields
+        if self.dynamic_key and not validate_cel_expression(self.dynamic_key):
+            raise ValueError(f"Invalid CEL expression: {self.dynamic_key}")
+
+        if not isinstance(self.units, int) and not validate_cel_expression(self.units):
+            raise ValueError(f"Invalid CEL expression: {self.units}")
+
+        if (
+            self.limit
+            and not isinstance(self.limit, int)
+            and not validate_cel_expression(self.limit)
+        ):
+            raise ValueError(f"Invalid CEL expression: {self.limit}")
+
+        if self.dynamic_key and not self.limit:
+            raise ValueError("CEL based keys requires limit to be set")
+
+        return self
+
+    def to_proto(self) -> CreateTaskRateLimit:
         key = self.static_key
         key_expression = self.dynamic_key
 
-        if self.key is not None:
-            DeprecationWarning(
-                "key is deprecated and will be removed in a future release, please use static_key instead"
-            )
-            key = self.key
-
-        if key_expression is not None:
-            if key is not None:
-                raise ValueError("Cannot have both static key and dynamic key set")
-
-            key = key_expression
-            if not validate_cel_expression(key_expression):
-                raise ValueError(f"Invalid CEL expression: {key_expression}")
-
-        # juggle the units and units_expr fields
-        units = None
-        units_expression = None
-        if isinstance(self.units, int):
-            units = self.units
-        else:
-            if not validate_cel_expression(self.units):
-                raise ValueError(f"Invalid CEL expression: {self.units}")
-            units_expression = self.units
-
-        # juggle the limit and limit_expr fields
-        limit_expression = None
-
-        if self.limit:
-            if isinstance(self.limit, int):
-                limit_expression = f"{self.limit}"
-            else:
-                if not validate_cel_expression(self.limit):
-                    raise ValueError(f"Invalid CEL expression: {self.limit}")
-                limit_expression = self.limit
-
-        if key_expression is not None and limit_expression is None:
-            raise ValueError("CEL based keys requires limit to be set")
+        key = self.static_key or self.dynamic_key
+
+        units = self.units if isinstance(self.units, int) else None
+        units_expression = None if isinstance(self.units, int) else self.units
+
+        limit_expression = None if not self.limit else str(self.limit)
 
-        self._req = CreateStepRateLimit(
+        return CreateTaskRateLimit(
             key=key,
             key_expr=key_expression,
             units=units,

hatchet_sdk/runnables/contextvars.py

@@ -0,0 +1,12 @@
+import asyncio
+from collections import Counter
+from contextvars import ContextVar
+
+ctx_workflow_run_id: ContextVar[str | None] = ContextVar(
+    "ctx_workflow_run_id", default=None
+)
+ctx_step_run_id: ContextVar[str | None] = ContextVar("ctx_step_run_id", default=None)
+ctx_worker_id: ContextVar[str | None] = ContextVar("ctx_worker_id", default=None)
+
+workflow_spawn_indices = Counter[str]()
+spawn_index_lock = asyncio.Lock()

hatchet_sdk/runnables/standalone.py

@@ -0,0 +1,194 @@
+import asyncio
+from datetime import datetime
+from typing import Any, Generic, cast, get_type_hints
+
+from google.protobuf import timestamp_pb2
+
+from hatchet_sdk.clients.admin import (
+    ScheduleTriggerWorkflowOptions,
+    TriggerWorkflowOptions,
+    WorkflowRunTriggerConfig,
+)
+from hatchet_sdk.clients.rest.models.cron_workflows import CronWorkflows
+from hatchet_sdk.contracts.workflows_pb2 import WorkflowVersion
+from hatchet_sdk.runnables.task import Task
+from hatchet_sdk.runnables.types import R, TWorkflowInput
+from hatchet_sdk.runnables.workflow import BaseWorkflow, Workflow
+from hatchet_sdk.utils.aio_utils import get_active_event_loop
+from hatchet_sdk.utils.typing import JSONSerializableMapping, is_basemodel_subclass
+from hatchet_sdk.workflow_run import WorkflowRunRef
+
+
+class TaskRunRef(Generic[TWorkflowInput, R]):
+    def __init__(
+        self,
+        standalone: "Standalone[TWorkflowInput, R]",
+        workflow_run_ref: WorkflowRunRef,
+    ):
+        self._s = standalone
+        self._wrr = workflow_run_ref
+
+    async def aio_result(self) -> R:
+        result = await self._wrr.workflow_listener.result(self._wrr.workflow_run_id)
+        return self._s._extract_result(result)
+
+    def result(self) -> R:
+        coro = self._wrr.workflow_listener.result(self._wrr.workflow_run_id)
+
+        loop = get_active_event_loop()
+
+        if loop is None:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                result = loop.run_until_complete(coro)
+            finally:
+                asyncio.set_event_loop(None)
+        else:
+            result = loop.run_until_complete(coro)
+
+        return self._s._extract_result(result)
+
+
+class Standalone(BaseWorkflow[TWorkflowInput], Generic[TWorkflowInput, R]):
+    def __init__(
+        self, workflow: Workflow[TWorkflowInput], task: Task[TWorkflowInput, R]
+    ) -> None:
+        super().__init__(config=workflow.config, client=workflow.client)
+
+        ## NOTE: This is a hack to assign the task back to the base workflow,
+        ## since the decorator to mutate the tasks is not being called.
+        self._default_tasks = [task]
+
+        self._workflow = workflow
+        self._task = task
+
+        return_type = get_type_hints(self._task.fn).get("return")
+
+        self._output_validator = (
+            return_type if is_basemodel_subclass(return_type) else None
+        )
+
+        self.config = self._workflow.config
+
+    def _extract_result(self, result: dict[str, Any]) -> R:
+        output = result.get(self._task.name)
+
+        if not self._output_validator:
+            return cast(R, output)
+
+        return cast(R, self._output_validator.model_validate(output))
+
+    def run(
+        self,
+        input: TWorkflowInput | None = None,
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> R:
+        return self._extract_result(self._workflow.run(input, options))
+
+    async def aio_run(
+        self,
+        input: TWorkflowInput | None = None,
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> R:
+        result = await self._workflow.aio_run(input, options)
+        return self._extract_result(result)
+
+    def run_no_wait(
+        self,
+        input: TWorkflowInput | None = None,
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> TaskRunRef[TWorkflowInput, R]:
+        ref = self._workflow.run_no_wait(input, options)
+
+        return TaskRunRef[TWorkflowInput, R](self, ref)
+
+    async def aio_run_no_wait(
+        self,
+        input: TWorkflowInput | None = None,
+        options: TriggerWorkflowOptions = TriggerWorkflowOptions(),
+    ) -> TaskRunRef[TWorkflowInput, R]:
+        ref = await self._workflow.aio_run_no_wait(input, options)
+
+        return TaskRunRef[TWorkflowInput, R](self, ref)
+
+    def run_many(self, workflows: list[WorkflowRunTriggerConfig]) -> list[R]:
+        return [
+            self._extract_result(result)
+            for result in self._workflow.run_many(workflows)
+        ]
+
+    async def aio_run_many(self, workflows: list[WorkflowRunTriggerConfig]) -> list[R]:
+        return [
+            self._extract_result(result)
+            for result in await self._workflow.aio_run_many(workflows)
+        ]
+
+    def run_many_no_wait(
+        self, workflows: list[WorkflowRunTriggerConfig]
+    ) -> list[TaskRunRef[TWorkflowInput, R]]:
+        refs = self._workflow.run_many_no_wait(workflows)
+
+        return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]
+
+    async def aio_run_many_no_wait(
+        self, workflows: list[WorkflowRunTriggerConfig]
+    ) -> list[TaskRunRef[TWorkflowInput, R]]:
+        refs = await self._workflow.aio_run_many_no_wait(workflows)
+
+        return [TaskRunRef[TWorkflowInput, R](self, ref) for ref in refs]
+
+    def schedule(
+        self,
+        schedules: list[datetime],
+        input: TWorkflowInput | None = None,
+        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
+    ) -> WorkflowVersion:
+        return self._workflow.schedule(
+            schedules=schedules,
+            input=input,
+            options=options,
+        )
+
+    async def aio_schedule(
+        self,
+        schedules: list[datetime | timestamp_pb2.Timestamp],
+        input: TWorkflowInput,
+        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
+    ) -> WorkflowVersion:
+        return await self._workflow.aio_schedule(
+            schedules=schedules,
+            input=input,
+            options=options,
+        )
+
+    def create_cron(
+        self,
+        cron_name: str,
+        expression: str,
+        input: TWorkflowInput,
+        additional_metadata: JSONSerializableMapping,
+    ) -> CronWorkflows:
+        return self._workflow.create_cron(
+            cron_name=cron_name,
+            expression=expression,
+            input=input,
+            additional_metadata=additional_metadata,
+        )
+
+    async def aio_create_cron(
+        self,
+        cron_name: str,
+        expression: str,
+        input: TWorkflowInput,
+        additional_metadata: JSONSerializableMapping,
+    ) -> CronWorkflows:
+        return await self._workflow.aio_create_cron(
+            cron_name=cron_name,
+            expression=expression,
+            input=input,
+            additional_metadata=additional_metadata,
+        )
+
+    def to_task(self) -> Task[TWorkflowInput, R]:
+        return self._task

hatchet_sdk/runnables/task.py

@@ -0,0 +1,144 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Generic,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from hatchet_sdk.context.context import Context, DurableContext
+from hatchet_sdk.contracts.v1.workflows_pb2 import (
+    CreateTaskRateLimit,
+    DesiredWorkerLabels,
+)
+from hatchet_sdk.runnables.types import (
+    DEFAULT_EXECUTION_TIMEOUT,
+    DEFAULT_SCHEDULE_TIMEOUT,
+    ConcurrencyExpression,
+    R,
+    StepType,
+    TWorkflowInput,
+    is_async_fn,
+    is_durable_sync_fn,
+    is_sync_fn,
+)
+from hatchet_sdk.utils.timedelta_to_expression import Duration
+from hatchet_sdk.waits import Condition, OrGroup
+
+if TYPE_CHECKING:
+    from hatchet_sdk.runnables.workflow import Workflow
+
+
+T = TypeVar("T")
+
+
+def fall_back_to_default(value: T, default: T, fallback_value: T) -> T:
+    ## If the value is not the default, it's set
+    if value != default:
+        return value
+
+    ## Otherwise, it's unset, so return the fallback value
+    return fallback_value
+
+
+class Task(Generic[TWorkflowInput, R]):
+    def __init__(
+        self,
+        _fn: Union[
+            Callable[[TWorkflowInput, Context], R]
+            | Callable[[TWorkflowInput, Context], Awaitable[R]],
+            Callable[[TWorkflowInput, DurableContext], R]
+            | Callable[[TWorkflowInput, DurableContext], Awaitable[R]],
+        ],
+        is_durable: bool,
+        type: StepType,
+        workflow: "Workflow[TWorkflowInput]",
+        name: str,
+        execution_timeout: Duration = DEFAULT_EXECUTION_TIMEOUT,
+        schedule_timeout: Duration = DEFAULT_SCHEDULE_TIMEOUT,
+        parents: "list[Task[TWorkflowInput, Any]]" = [],
+        retries: int = 0,
+        rate_limits: list[CreateTaskRateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabels] = {},
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+        concurrency: list[ConcurrencyExpression] = [],
+        wait_for: list[Condition | OrGroup] = [],
+        skip_if: list[Condition | OrGroup] = [],
+        cancel_if: list[Condition | OrGroup] = [],
+    ) -> None:
+        self.is_durable = is_durable
+
+        self.fn = _fn
+        self.is_async_function = is_async_fn(self.fn)  # type: ignore
+
+        self.workflow = workflow
+
+        self.type = type
+        self.execution_timeout = fall_back_to_default(
+            execution_timeout, DEFAULT_EXECUTION_TIMEOUT, DEFAULT_EXECUTION_TIMEOUT
+        )
+        self.schedule_timeout = fall_back_to_default(
+            schedule_timeout, DEFAULT_SCHEDULE_TIMEOUT, DEFAULT_SCHEDULE_TIMEOUT
+        )
+        self.name = name
+        self.parents = parents
+        self.retries = retries
+        self.rate_limits = rate_limits
+        self.desired_worker_labels = desired_worker_labels
+        self.backoff_factor = backoff_factor
+        self.backoff_max_seconds = backoff_max_seconds
+        self.concurrency = concurrency
+
+        self.wait_for = self._flatten_conditions(wait_for)
+        self.skip_if = self._flatten_conditions(skip_if)
+        self.cancel_if = self._flatten_conditions(cancel_if)
+
+    def _flatten_conditions(
+        self, conditions: list[Condition | OrGroup]
+    ) -> list[Condition]:
+        flattened: list[Condition] = []
+
+        for condition in conditions:
+            if isinstance(condition, OrGroup):
+                for or_condition in condition.conditions:
+                    or_condition.base.or_group_id = condition.or_group_id
+
+                flattened.extend(condition.conditions)
+            else:
+                flattened.append(condition)
+
+        return flattened
+
+    def call(self, ctx: Context | DurableContext) -> R:
+        if self.is_async_function:
+            raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
+
+        workflow_input = self.workflow._get_workflow_input(ctx)
+
+        if self.is_durable:
+            fn = cast(Callable[[TWorkflowInput, DurableContext], R], self.fn)
+            if is_durable_sync_fn(fn):
+                return fn(workflow_input, cast(DurableContext, ctx))
+        else:
+            fn = cast(Callable[[TWorkflowInput, Context], R], self.fn)
+            if is_sync_fn(fn):
+                return fn(workflow_input, cast(Context, ctx))
+
+        raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
+
+    async def aio_call(self, ctx: Context | DurableContext) -> R:
+        if not self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not an async function. Use `call` instead."
+            )
+
+        workflow_input = self.workflow._get_workflow_input(ctx)
+
+        if is_async_fn(self.fn):  # type: ignore
+            return await self.fn(workflow_input, cast(Context, ctx))  # type: ignore
+
+        raise TypeError(f"{self.name} is not an async function. Use `call` instead.")