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

hatchet_sdk/hatchet.py

@@ -1,197 +1,46 @@
 import asyncio
 import logging
-from typing import Any, Callable, Optional, ParamSpec, Type, TypeVar, Union
-
-from pydantic import BaseModel
-from typing_extensions import deprecated
+from typing import Any, Callable, Type, TypeVar, cast, overload
 
+from hatchet_sdk.client import Client, new_client, new_client_raw
+from hatchet_sdk.clients.admin import AdminClient
+from hatchet_sdk.clients.dispatcher.dispatcher import DispatcherClient
+from hatchet_sdk.clients.events import EventClient
 from hatchet_sdk.clients.rest_client import RestApi
+from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
+from hatchet_sdk.config import ClientConfig
 from hatchet_sdk.context.context import Context
-from hatchet_sdk.contracts.workflows_pb2 import (
-    ConcurrencyLimitStrategy,
-    CreateStepRateLimit,
-    DesiredWorkerLabels,
-    StickyStrategy,
-)
+from hatchet_sdk.contracts.workflows_pb2 import DesiredWorkerLabels
 from hatchet_sdk.features.cron import CronClient
 from hatchet_sdk.features.scheduled import ScheduledClient
 from hatchet_sdk.labels import DesiredWorkerLabel
-from hatchet_sdk.loader import ClientConfig, ConfigLoader
+from hatchet_sdk.logger import logger
 from hatchet_sdk.rate_limit import RateLimit
-from hatchet_sdk.v2.callable import HatchetCallable
-
-from .client import Client, new_client, new_client_raw
-from .clients.admin import AdminClient
-from .clients.dispatcher.dispatcher import DispatcherClient
-from .clients.events import EventClient
-from .clients.run_event_listener import RunEventListenerClient
-from .logger import logger
-from .worker.worker import Worker
-from .workflow import (
+from hatchet_sdk.worker.worker import Worker
+from hatchet_sdk.workflow import (
     ConcurrencyExpression,
-    WorkflowInterface,
-    WorkflowMeta,
-    WorkflowStepProtocol,
+    EmptyModel,
+    Step,
+    StepType,
+    StickyStrategy,
+    Task,
+    TWorkflowInput,
+    WorkflowConfig,
+    WorkflowDeclaration,
 )
 
-T = TypeVar("T", bound=BaseModel)
 R = TypeVar("R")
-P = ParamSpec("P")
-
-TWorkflow = TypeVar("TWorkflow", bound=object)
-
-
-def workflow(
-    name: str = "",
-    on_events: list[str] | None = None,
-    on_crons: list[str] | None = None,
-    version: str = "",
-    timeout: str = "60m",
-    schedule_timeout: str = "5m",
-    sticky: Union[StickyStrategy.Value, None] = None,  # type: ignore[name-defined]
-    default_priority: int | None = None,
-    concurrency: ConcurrencyExpression | None = None,
-    input_validator: Type[T] | None = None,
-) -> Callable[[Type[TWorkflow]], WorkflowMeta]:
-    on_events = on_events or []
-    on_crons = on_crons or []
-
-    def inner(cls: Type[TWorkflow]) -> WorkflowMeta:
-        nonlocal name
-        name = name or str(cls.__name__)
-
-        setattr(cls, "on_events", on_events)
-        setattr(cls, "on_crons", on_crons)
-        setattr(cls, "name", name)
-        setattr(cls, "version", version)
-        setattr(cls, "timeout", timeout)
-        setattr(cls, "schedule_timeout", schedule_timeout)
-        setattr(cls, "sticky", sticky)
-        setattr(cls, "default_priority", default_priority)
-        setattr(cls, "concurrency_expression", concurrency)
-
-        # Define a new class with the same name and bases as the original, but
-        # with WorkflowMeta as its metaclass
-
-        ## TODO: Figure out how to type this metaclass correctly
-        setattr(cls, "input_validator", input_validator)
-
-        return WorkflowMeta(name, cls.__bases__, dict(cls.__dict__))
-
-    return inner
-
-
-def step(
-    name: str = "",
-    timeout: str = "",
-    parents: list[str] | None = None,
-    retries: int = 0,
-    rate_limits: list[RateLimit] | None = None,
-    desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
-    backoff_factor: float | None = None,
-    backoff_max_seconds: int | None = None,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    parents = parents or []
-
-    def inner(func: Callable[P, R]) -> Callable[P, R]:
-        limits = None
-        if rate_limits:
-            limits = [rate_limit._req for rate_limit in rate_limits or []]
-
-        setattr(func, "_step_name", name.lower() or str(func.__name__).lower())
-        setattr(func, "_step_parents", parents)
-        setattr(func, "_step_timeout", timeout)
-        setattr(func, "_step_retries", retries)
-        setattr(func, "_step_rate_limits", limits)
-        setattr(func, "_step_backoff_factor", backoff_factor)
-        setattr(func, "_step_backoff_max_seconds", backoff_max_seconds)
-
-        def create_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels:
-            value = d["value"] if "value" in d else None
-            return DesiredWorkerLabels(
-                strValue=str(value) if not isinstance(value, int) else None,
-                intValue=value if isinstance(value, int) else None,
-                required=d["required"] if "required" in d else None,  # type: ignore[arg-type]
-                weight=d["weight"] if "weight" in d else None,
-                comparator=d["comparator"] if "comparator" in d else None,  # type: ignore[arg-type]
-            )
-
-        setattr(
-            func,
-            "_step_desired_worker_labels",
-            {key: create_label(d) for key, d in desired_worker_labels.items()},
-        )
-
-        return func
-
-    return inner
-
-
-def on_failure_step(
-    name: str = "",
-    timeout: str = "",
-    retries: int = 0,
-    rate_limits: list[RateLimit] | None = None,
-    backoff_factor: float | None = None,
-    backoff_max_seconds: int | None = None,
-) -> Callable[..., Any]:
-    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
-        limits = None
-        if rate_limits:
-            limits = [
-                CreateStepRateLimit(key=rate_limit.static_key, units=rate_limit.units)  # type: ignore[arg-type]
-                for rate_limit in rate_limits or []
-            ]
-
-        setattr(
-            func, "_on_failure_step_name", name.lower() or str(func.__name__).lower()
-        )
-        setattr(func, "_on_failure_step_timeout", timeout)
-        setattr(func, "_on_failure_step_retries", retries)
-        setattr(func, "_on_failure_step_rate_limits", limits)
-        setattr(func, "_on_failure_step_backoff_factor", backoff_factor)
-        setattr(func, "_on_failure_step_backoff_max_seconds", backoff_max_seconds)
-
-        return func
-
-    return inner
-
-
-def concurrency(
-    name: str = "",
-    max_runs: int = 1,
-    limit_strategy: ConcurrencyLimitStrategy = ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS,
-) -> Callable[..., Any]:
-    def inner(func: Callable[[Context], Any]) -> Callable[[Context], Any]:
-        setattr(
-            func,
-            "_concurrency_fn_name",
-            name.lower() or str(func.__name__).lower(),
-        )
-        setattr(func, "_concurrency_max_runs", max_runs)
-        setattr(func, "_concurrency_limit_strategy", limit_strategy)
-
-        return func
-
-    return inner
-
-
-class HatchetRest:
-    """
-    Main client for interacting with the Hatchet API.
-
-    This class provides access to various client interfaces and utility methods
-    for working with Hatchet via the REST API,
-
-    Attributes:
-        rest (RestApi): Interface for REST API operations.
-    """
 
-    rest: RestApi
 
-    def __init__(self, config: ClientConfig = ClientConfig()):
-        _config: ClientConfig = ConfigLoader(".").load_client_config(config)
-        self.rest = RestApi(_config.server_url, _config.token, _config.tenant_id)
+def transform_desired_worker_label(d: DesiredWorkerLabel) -> DesiredWorkerLabels:
+    value = d.value
+    return DesiredWorkerLabels(
+        strValue=value if not isinstance(value, int) else None,
+        intValue=value if isinstance(value, int) else None,
+        required=d.required,
+        weight=d.weight,
+        comparator=d.comparator,  # type: ignore[arg-type]
+    )
 
 
 class Hatchet:
@@ -227,7 +76,7 @@ class Hatchet:
     def __init__(
         self,
         debug: bool = False,
-        client: Optional[Client] = None,
+        client: Client | None = None,
         config: ClientConfig = ClientConfig(),
     ):
         """
@@ -238,24 +87,14 @@ class Hatchet:
             client (Optional[Client], optional): A pre-configured Client instance. Defaults to None.
             config (ClientConfig, optional): Configuration for creating a new Client. Defaults to ClientConfig().
         """
-        if client is not None:
-            self._client = client
-        else:
-            self._client = new_client(config, debug)
 
         if debug:
             logger.setLevel(logging.DEBUG)
 
+        self._client = client if client else new_client(config, debug)
         self.cron = CronClient(self._client)
         self.scheduled = ScheduledClient(self._client)
 
-    @property
-    @deprecated(
-        "Direct access to client is deprecated and will be removed in a future version. Use specific client properties (Hatchet.admin, Hatchet.dispatcher, Hatchet.event, Hatchet.rest) instead. [0.32.0]",
-    )
-    def client(self) -> Client:
-        return self._client
-
     @property
     def admin(self) -> AdminClient:
         return self._client.admin
@@ -284,13 +123,108 @@ class Hatchet:
     def tenant_id(self) -> str:
         return self._client.config.tenant_id
 
-    concurrency = staticmethod(concurrency)
+    def step(
+        self,
+        name: str = "",
+        timeout: str = "60m",
+        parents: list[str] = [],
+        retries: int = 0,
+        rate_limits: list[RateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+    ) -> Callable[[Callable[[Any, Context], R]], Step[R]]:
+        def inner(func: Callable[[Any, Context], R]) -> Step[R]:
+            return Step(
+                fn=func,
+                type=StepType.DEFAULT,
+                name=name.lower() or str(func.__name__).lower(),
+                timeout=timeout,
+                parents=parents,
+                retries=retries,
+                rate_limits=[r for rate_limit in rate_limits if (r := rate_limit._req)],
+                desired_worker_labels={
+                    key: transform_desired_worker_label(d)
+                    for key, d in desired_worker_labels.items()
+                },
+                backoff_factor=backoff_factor,
+                backoff_max_seconds=backoff_max_seconds,
+            )
 
-    workflow = staticmethod(workflow)
+        return inner
 
-    step = staticmethod(step)
+    def on_failure_step(
+        self,
+        name: str = "",
+        timeout: str = "60m",
+        parents: list[str] = [],
+        retries: int = 0,
+        rate_limits: list[RateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+    ) -> Callable[[Callable[[Any, Context], R]], Step[R]]:
+        def inner(func: Callable[[Any, Context], R]) -> Step[R]:
+            return Step(
+                fn=func,
+                type=StepType.ON_FAILURE,
+                name=name.lower() or str(func.__name__).lower(),
+                timeout=timeout,
+                parents=parents,
+                retries=retries,
+                rate_limits=[r for rate_limit in rate_limits if (r := rate_limit._req)],
+                desired_worker_labels={
+                    key: transform_desired_worker_label(d)
+                    for key, d in desired_worker_labels.items()
+                },
+                backoff_factor=backoff_factor,
+                backoff_max_seconds=backoff_max_seconds,
+            )
 
-    on_failure_step = staticmethod(on_failure_step)
+        return inner
+
+    def task(
+        self,
+        name: str = "",
+        on_events: list[str] = [],
+        on_crons: list[str] = [],
+        version: str = "",
+        timeout: str = "60m",
+        schedule_timeout: str = "5m",
+        sticky: StickyStrategy | None = None,
+        retries: int = 0,
+        rate_limits: list[RateLimit] = [],
+        desired_worker_labels: dict[str, DesiredWorkerLabel] = {},
+        concurrency: ConcurrencyExpression | None = None,
+        on_failure: Task[Any, Any] | None = None,
+        default_priority: int = 1,
+        input_validator: Type[TWorkflowInput] | None = None,
+        backoff_factor: float | None = None,
+        backoff_max_seconds: int | None = None,
+    ) -> Callable[[Callable[[Context], R]], Task[R, TWorkflowInput]]:
+        def inner(func: Callable[[Context], R]) -> Task[R, TWorkflowInput]:
+            return Task[R, TWorkflowInput](
+                func,
+                hatchet=self,
+                name=name,
+                on_events=on_events,
+                on_crons=on_crons,
+                version=version,
+                timeout=timeout,
+                schedule_timeout=schedule_timeout,
+                sticky=sticky,
+                retries=retries,
+                rate_limits=rate_limits,
+                desired_worker_labels=desired_worker_labels,
+                concurrency=concurrency,
+                on_failure=on_failure,
+                default_priority=default_priority,
+                input_validator=input_validator,
+                backoff_factor=backoff_factor,
+                backoff_max_seconds=backoff_max_seconds,
+            )
+
+        return inner
 
     def worker(
         self, name: str, max_runs: int | None = None, labels: dict[str, str | int] = {}
@@ -308,3 +242,66 @@ class Hatchet:
             debug=self._client.debug,
             owned_loop=loop is None,
         )
+
+    @overload
+    def declare_workflow(
+        self,
+        *,
+        name: str = "",
+        on_events: list[str] = [],
+        on_crons: list[str] = [],
+        version: str = "",
+        timeout: str = "60m",
+        schedule_timeout: str = "5m",
+        sticky: StickyStrategy | None = None,
+        default_priority: int = 1,
+        concurrency: ConcurrencyExpression | None = None,
+        input_validator: None = None,
+    ) -> WorkflowDeclaration[EmptyModel]: ...
+
+    @overload
+    def declare_workflow(
+        self,
+        *,
+        name: str = "",
+        on_events: list[str] = [],
+        on_crons: list[str] = [],
+        version: str = "",
+        timeout: str = "60m",
+        schedule_timeout: str = "5m",
+        sticky: StickyStrategy | None = None,
+        default_priority: int = 1,
+        concurrency: ConcurrencyExpression | None = None,
+        input_validator: Type[TWorkflowInput],
+    ) -> WorkflowDeclaration[TWorkflowInput]: ...
+
+    def declare_workflow(
+        self,
+        *,
+        name: str = "",
+        on_events: list[str] = [],
+        on_crons: list[str] = [],
+        version: str = "",
+        timeout: str = "60m",
+        schedule_timeout: str = "5m",
+        sticky: StickyStrategy | None = None,
+        default_priority: int = 1,
+        concurrency: ConcurrencyExpression | None = None,
+        input_validator: Type[TWorkflowInput] | None = None,
+    ) -> WorkflowDeclaration[EmptyModel] | WorkflowDeclaration[TWorkflowInput]:
+        return WorkflowDeclaration[TWorkflowInput](
+            WorkflowConfig(
+                name=name,
+                on_events=on_events,
+                on_crons=on_crons,
+                version=version,
+                timeout=timeout,
+                schedule_timeout=schedule_timeout,
+                sticky=sticky,
+                default_priority=default_priority,
+                concurrency=concurrency,
+                input_validator=input_validator
+                or cast(Type[TWorkflowInput], EmptyModel),
+            ),
+            self,
+        )

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]
@@ -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,
         )
 
@@ -391,6 +395,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,7 +1,7 @@
 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 hatchet_sdk.contracts.workflows_pb2 import CreateStepRateLimit
 
@@ -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"
@@ -40,7 +40,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,29 +60,21 @@ 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
+    _req: CreateStepRateLimit | None = None
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         # juggle the key and key_expr fields
         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")

hatchet_sdk/token.py

@@ -1,20 +1,25 @@
 import base64
-import json
 
+from pydantic import BaseModel
 
-def get_tenant_id_from_jwt(token: str) -> str:
-    claims = extract_claims_from_jwt(token)
 
-    return claims.get("sub")
+class Claims(BaseModel):
+    sub: str
+    server_url: str
+    grpc_broadcast_address: str
+
+
+def get_tenant_id_from_jwt(token: str) -> str:
+    return extract_claims_from_jwt(token).sub
 
 
-def get_addresses_from_jwt(token: str) -> (str, str):
+def get_addresses_from_jwt(token: str) -> tuple[str, str]:
     claims = extract_claims_from_jwt(token)
 
-    return claims.get("server_url"), claims.get("grpc_broadcast_address")
+    return claims.server_url, claims.grpc_broadcast_address
 
 
-def extract_claims_from_jwt(token: str):
+def extract_claims_from_jwt(token: str) -> Claims:
     parts = token.split(".")
     if len(parts) != 3:
         raise ValueError("Invalid token format")
@@ -22,6 +27,5 @@ def extract_claims_from_jwt(token: str):
     claims_part = parts[1]
     claims_part += "=" * ((4 - len(claims_part) % 4) % 4)  # Padding for base64 decoding
     claims_data = base64.urlsafe_b64decode(claims_part)
-    claims = json.loads(claims_data)
 
-    return claims
+    return Claims.model_validate_json(claims_data)

hatchet_sdk/utils/aio_utils.py

@@ -2,7 +2,6 @@ import asyncio
 import inspect
 from concurrent.futures import Executor
 from functools import partial, wraps
-from threading import Thread
 from typing import Any
 
 
@@ -81,45 +80,6 @@ def sync_to_async(func: Any) -> Any:
     return run
 
 
-class EventLoopThread:
-    """A class that manages an asyncio event loop running in a separate thread."""
-
-    def __init__(self) -> None:
-        """
-        Initializes the EventLoopThread by creating an event loop
-        and setting up a thread to run the loop.
-        """
-        self.loop = asyncio.new_event_loop()
-        self.thread = Thread(target=self.run_loop_in_thread, args=(self.loop,))
-
-    def __enter__(self, *a, **kw) -> asyncio.AbstractEventLoop:
-        """
-        Starts the thread running the event loop when entering the context.
-
-        Returns:
-            asyncio.AbstractEventLoop: The event loop running in the separate thread.
-        """
-        self.thread.start()
-        return self.loop
-
-    def __exit__(self, *a, **kw) -> None:
-        """
-        Stops the event loop and joins the thread when exiting the context.
-        """
-        self.loop.call_soon_threadsafe(self.loop.stop)
-        self.thread.join()
-
-    def run_loop_in_thread(self, loop: asyncio.AbstractEventLoop) -> None:
-        """
-        Sets the event loop for the current thread and runs it forever.
-
-        Args:
-            loop (asyncio.AbstractEventLoop): The event loop to run.
-        """
-        asyncio.set_event_loop(loop)
-        loop.run_forever()
-
-
 def get_active_event_loop() -> asyncio.AbstractEventLoop | None:
     """
     Get the active event loop.