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

hatchet_sdk/v0/opentelemetry/instrumentor.py

@@ -0,0 +1,396 @@
+from importlib.metadata import version
+from typing import Any, Callable, Collection, Coroutine
+
+try:
+    from opentelemetry.context import Context
+    from opentelemetry.instrumentation.instrumentor import (  # type: ignore[attr-defined]
+        BaseInstrumentor,
+    )
+    from opentelemetry.instrumentation.utils import unwrap
+    from opentelemetry.metrics import MeterProvider, NoOpMeterProvider, get_meter
+    from opentelemetry.trace import (
+        NoOpTracerProvider,
+        StatusCode,
+        TracerProvider,
+        get_tracer,
+        get_tracer_provider,
+    )
+    from opentelemetry.trace.propagation.tracecontext import (
+        TraceContextTextMapPropagator,
+    )
+    from wrapt import wrap_function_wrapper  # type: ignore[import-untyped]
+except (RuntimeError, ImportError, ModuleNotFoundError):
+    raise ModuleNotFoundError(
+        "To use the HatchetInstrumentor, you must install Hatchet's `otel` extra using (e.g.) `pip install hatchet-sdk[otel]`"
+    )
+
+import hatchet_sdk
+from hatchet_sdk.v0.clients.admin import (
+    AdminClient,
+    TriggerWorkflowOptions,
+    WorkflowRunDict,
+)
+from hatchet_sdk.v0.clients.dispatcher.action_listener import Action
+from hatchet_sdk.v0.clients.events import (
+    BulkPushEventWithMetadata,
+    EventClient,
+    PushEventOptions,
+)
+from hatchet_sdk.v0.contracts.events_pb2 import Event
+from hatchet_sdk.v0.worker.runner.runner import Runner
+from hatchet_sdk.v0.workflow_run import WorkflowRunRef
+
+hatchet_sdk_version = version("hatchet-sdk")
+
+InstrumentKwargs = TracerProvider | MeterProvider | None
+
+OTEL_TRACEPARENT_KEY = "traceparent"
+
+
+def create_traceparent() -> str | None:
+    """
+    Creates and returns a W3C traceparent header value using OpenTelemetry's context propagation.
+
+    The traceparent header is used to propagate context information across service boundaries
+    in distributed tracing systems. It follows the W3C Trace Context specification.
+
+    :returns: A W3C-formatted traceparent header value if successful, None if the context
+                    injection fails or no active span exists.\n
+                    Example: `00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01`
+    :rtype: str | None:
+    """
+
+    carrier: dict[str, str] = {}
+    TraceContextTextMapPropagator().inject(carrier)
+
+    return carrier.get("traceparent")
+
+
+def parse_carrier_from_metadata(metadata: dict[str, str] | None) -> Context | None:
+    """
+    Parses OpenTelemetry trace context from a metadata dictionary.
+
+    Extracts the trace context from metadata using the W3C Trace Context format,
+    specifically looking for the `traceparent` header.
+
+    :param metadata: A dictionary containing metadata key-value pairs,
+                     potentially including the `traceparent` header. Can be None.
+    :type metadata: dict[str, str] | None
+    :returns: The extracted OpenTelemetry Context object if a valid `traceparent`
+              is found in the metadata, otherwise None.
+    :rtype: Context | None
+
+    :Example:
+
+    >>> metadata = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"}
+    >>> context = parse_carrier_from_metadata(metadata)
+    """
+
+    if not metadata:
+        return None
+
+    traceparent = metadata.get(OTEL_TRACEPARENT_KEY)
+
+    if not traceparent:
+        return None
+
+    return TraceContextTextMapPropagator().extract({OTEL_TRACEPARENT_KEY: traceparent})
+
+
+def inject_traceparent_into_metadata(
+    metadata: dict[str, str], traceparent: str | None = None
+) -> dict[str, str]:
+    """
+    Injects OpenTelemetry `traceparent` into a metadata dictionary.
+
+    Takes a metadata dictionary and an optional `traceparent` string,
+    returning a new metadata dictionary with the `traceparent` added under the
+    `OTEL_TRACEPARENT_KEY`. If no `traceparent` is provided, it attempts to create one.
+
+    :param metadata: The metadata dictionary to inject the `traceparent` into.
+    :type metadata: dict[str, str]
+    :param traceparent: The `traceparent` string to inject. If None, attempts to use
+                        the current span.
+    :type traceparent: str | None, optional
+    :returns: A new metadata dictionary containing the original metadata plus
+              the injected `traceparent`, if one was available or could be created.
+    :rtype: dict[str, str]
+
+    :Example:
+
+    >>> metadata = {"key": "value"}
+    >>> new_metadata = inject_traceparent(metadata, "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01")
+    >>> print(new_metadata)
+    {"key": "value", "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
+    """
+
+    if not traceparent:
+        traceparent = create_traceparent()
+
+    if not traceparent:
+        return metadata
+
+    return {
+        **metadata,
+        OTEL_TRACEPARENT_KEY: traceparent,
+    }
+
+
+class HatchetInstrumentor(BaseInstrumentor):  # type: ignore[misc]
+    def __init__(
+        self,
+        tracer_provider: TracerProvider | None = None,
+        meter_provider: MeterProvider | None = None,
+    ):
+        """
+        Hatchet OpenTelemetry instrumentor.
+
+        The instrumentor provides an OpenTelemetry integration for Hatchet by setting up
+        tracing and metrics collection.
+
+        :param tracer_provider: TracerProvider | None: The OpenTelemetry TracerProvider to use.
+                If not provided, the global tracer provider will be used.
+        :param meter_provider: MeterProvider | None: The OpenTelemetry MeterProvider to use.
+                If not provided, a no-op meter provider will be used.
+        """
+
+        self.tracer_provider = tracer_provider or get_tracer_provider()
+        self.meter_provider = meter_provider or NoOpMeterProvider()
+
+        super().__init__()
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return tuple()
+
+    def _instrument(self, **kwargs: InstrumentKwargs) -> None:
+        self._tracer = get_tracer(__name__, hatchet_sdk_version, self.tracer_provider)
+        self._meter = get_meter(__name__, hatchet_sdk_version, self.meter_provider)
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "worker.runner.runner.Runner.handle_start_step_run",
+            self._wrap_handle_start_step_run,
+        )
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "worker.runner.runner.Runner.handle_start_group_key_run",
+            self._wrap_handle_get_group_key_run,
+        )
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "worker.runner.runner.Runner.handle_cancel_action",
+            self._wrap_handle_cancel_action,
+        )
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "clients.events.EventClient.push",
+            self._wrap_push_event,
+        )
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "clients.events.EventClient.bulk_push",
+            self._wrap_bulk_push_event,
+        )
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "clients.admin.AdminClient.run_workflow",
+            self._wrap_run_workflow,
+        )
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "clients.admin.AdminClientAioImpl.run_workflow",
+            self._wrap_async_run_workflow,
+        )
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "clients.admin.AdminClient.run_workflows",
+            self._wrap_run_workflows,
+        )
+
+        wrap_function_wrapper(
+            hatchet_sdk,
+            "clients.admin.AdminClientAioImpl.run_workflows",
+            self._wrap_async_run_workflows,
+        )
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    async def _wrap_handle_start_step_run(
+        self,
+        wrapped: Callable[[Action], Coroutine[None, None, Exception | None]],
+        instance: Runner,
+        args: tuple[Action],
+        kwargs: Any,
+    ) -> Exception | None:
+        action = args[0]
+        traceparent = parse_carrier_from_metadata(action.additional_metadata)
+
+        with self._tracer.start_as_current_span(
+            "hatchet.start_step_run",
+            attributes=action.otel_attributes,
+            context=traceparent,
+        ) as span:
+            result = await wrapped(*args, **kwargs)
+
+            if isinstance(result, Exception):
+                span.set_status(StatusCode.ERROR, str(result))
+
+            return result
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    async def _wrap_handle_get_group_key_run(
+        self,
+        wrapped: Callable[[Action], Coroutine[None, None, Exception | None]],
+        instance: Runner,
+        args: tuple[Action],
+        kwargs: Any,
+    ) -> Exception | None:
+        action = args[0]
+
+        with self._tracer.start_as_current_span(
+            "hatchet.get_group_key_run",
+            attributes=action.otel_attributes,
+        ) as span:
+            result = await wrapped(*args, **kwargs)
+
+            if isinstance(result, Exception):
+                span.set_status(StatusCode.ERROR, str(result))
+
+            return result
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    async def _wrap_handle_cancel_action(
+        self,
+        wrapped: Callable[[str], Coroutine[None, None, Exception | None]],
+        instance: Runner,
+        args: tuple[str],
+        kwargs: Any,
+    ) -> Exception | None:
+        step_run_id = args[0]
+
+        with self._tracer.start_as_current_span(
+            "hatchet.cancel_step_run",
+            attributes={
+                "hatchet.step_run_id": step_run_id,
+            },
+        ):
+            return await wrapped(*args, **kwargs)
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    def _wrap_push_event(
+        self,
+        wrapped: Callable[[str, dict[str, Any], PushEventOptions | None], Event],
+        instance: EventClient,
+        args: tuple[
+            str,
+            dict[str, Any],
+            PushEventOptions | None,
+        ],
+        kwargs: dict[str, str | dict[str, Any] | PushEventOptions | None],
+    ) -> Event:
+        with self._tracer.start_as_current_span(
+            "hatchet.push_event",
+        ):
+            return wrapped(*args, **kwargs)
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    def _wrap_bulk_push_event(
+        self,
+        wrapped: Callable[
+            [list[BulkPushEventWithMetadata], PushEventOptions | None], list[Event]
+        ],
+        instance: EventClient,
+        args: tuple[
+            list[BulkPushEventWithMetadata],
+            PushEventOptions | None,
+        ],
+        kwargs: dict[str, list[BulkPushEventWithMetadata] | PushEventOptions | None],
+    ) -> list[Event]:
+        with self._tracer.start_as_current_span(
+            "hatchet.bulk_push_event",
+        ):
+            return wrapped(*args, **kwargs)
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    def _wrap_run_workflow(
+        self,
+        wrapped: Callable[[str, Any, TriggerWorkflowOptions | None], WorkflowRunRef],
+        instance: AdminClient,
+        args: tuple[str, Any, TriggerWorkflowOptions | None],
+        kwargs: dict[str, str | Any | TriggerWorkflowOptions | None],
+    ) -> WorkflowRunRef:
+        with self._tracer.start_as_current_span(
+            "hatchet.run_workflow",
+        ):
+            return wrapped(*args, **kwargs)
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    async def _wrap_async_run_workflow(
+        self,
+        wrapped: Callable[
+            [str, Any, TriggerWorkflowOptions | None],
+            Coroutine[None, None, WorkflowRunRef],
+        ],
+        instance: AdminClient,
+        args: tuple[str, Any, TriggerWorkflowOptions | None],
+        kwargs: dict[str, str | Any | TriggerWorkflowOptions | None],
+    ) -> WorkflowRunRef:
+        with self._tracer.start_as_current_span(
+            "hatchet.run_workflow",
+        ):
+            return await wrapped(*args, **kwargs)
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    def _wrap_run_workflows(
+        self,
+        wrapped: Callable[
+            [list[WorkflowRunDict], TriggerWorkflowOptions | None], list[WorkflowRunRef]
+        ],
+        instance: AdminClient,
+        args: tuple[
+            list[WorkflowRunDict],
+            TriggerWorkflowOptions | None,
+        ],
+        kwargs: dict[str, list[WorkflowRunDict] | TriggerWorkflowOptions | None],
+    ) -> list[WorkflowRunRef]:
+        with self._tracer.start_as_current_span(
+            "hatchet.run_workflows",
+        ):
+            return wrapped(*args, **kwargs)
+
+    ## IMPORTANT: Keep these types in sync with the wrapped method's signature
+    async def _wrap_async_run_workflows(
+        self,
+        wrapped: Callable[
+            [list[WorkflowRunDict], TriggerWorkflowOptions | None],
+            Coroutine[None, None, list[WorkflowRunRef]],
+        ],
+        instance: AdminClient,
+        args: tuple[
+            list[WorkflowRunDict],
+            TriggerWorkflowOptions | None,
+        ],
+        kwargs: dict[str, list[WorkflowRunDict] | TriggerWorkflowOptions | None],
+    ) -> list[WorkflowRunRef]:
+        with self._tracer.start_as_current_span(
+            "hatchet.run_workflows",
+        ):
+            return await wrapped(*args, **kwargs)
+
+    def _uninstrument(self, **kwargs: InstrumentKwargs) -> None:
+        self.tracer_provider = NoOpTracerProvider()
+        self.meter_provider = NoOpMeterProvider()
+
+        unwrap(hatchet_sdk, "worker.runner.runner.Runner.handle_start_step_run")
+        unwrap(hatchet_sdk, "worker.runner.runner.Runner.handle_start_group_key_run")
+        unwrap(hatchet_sdk, "worker.runner.runner.Runner.handle_cancel_action")
+        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.run_workflows")
+        unwrap(hatchet_sdk, "clients.admin.AdminClientAioImpl.run_workflows")

hatchet_sdk/v0/rate_limit.py

@@ -0,0 +1,126 @@
+from dataclasses import dataclass
+from typing import Union
+
+from celpy import CELEvalError, Environment
+
+from hatchet_sdk.v0.contracts.workflows_pb2 import CreateStepRateLimit
+
+
+def validate_cel_expression(expr: str) -> bool:
+    env = Environment()
+    try:
+        env.compile(expr)
+        return True
+    except CELEvalError:
+        return False
+
+
+class RateLimitDuration:
+    SECOND = "SECOND"
+    MINUTE = "MINUTE"
+    HOUR = "HOUR"
+    DAY = "DAY"
+    WEEK = "WEEK"
+    MONTH = "MONTH"
+    YEAR = "YEAR"
+
+
+@dataclass
+class RateLimit:
+    """
+    Represents a rate limit configuration for a step in a workflow.
+
+    This class allows for both static and dynamic rate limiting based on various parameters.
+    It supports both simple integer values and Common Expression Language (CEL) expressions
+    for dynamic evaluation.
+
+    Attributes:
+        static_key (str, optional): A static key for rate limiting.
+        dynamic_key (str, optional): A CEL expression for dynamic key evaluation.
+        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:
+           rate_limit = RateLimit(static_key="external-api", units=100)
+           > NOTE: if you want to use a static key, you must first put the rate limit: hatchet.admin.put_rate_limit("external-api", 200, RateLimitDuration.SECOND)
+
+        2. Dynamic rate limit with CEL expressions:
+           rate_limit = RateLimit(
+               dynamic_key="input.user_id",
+               units="input.units",
+               limit="input.limit * input.user_tier"
+           )
+
+    Note:
+        - Either static_key or dynamic_key must be set, but not both.
+        - When using dynamic_key, limit must also be set.
+        - CEL expressions are validated upon instantiation.
+
+    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
+    duration: RateLimitDuration = RateLimitDuration.MINUTE
+
+    _req: CreateStepRateLimit = None
+
+    def __post_init__(self):
+        # 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")
+
+            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")
+
+        self._req = CreateStepRateLimit(
+            key=key,
+            key_expr=key_expression,
+            units=units,
+            units_expr=units_expression,
+            limit_values_expr=limit_expression,
+            duration=self.duration,
+        )

hatchet_sdk/v0/token.py

@@ -0,0 +1,27 @@
+import base64
+import json
+
+
+def get_tenant_id_from_jwt(token: str) -> str:
+    claims = extract_claims_from_jwt(token)
+
+    return claims.get("sub")
+
+
+def get_addresses_from_jwt(token: str) -> (str, str):
+    claims = extract_claims_from_jwt(token)
+
+    return claims.get("server_url"), claims.get("grpc_broadcast_address")
+
+
+def extract_claims_from_jwt(token: str):
+    parts = token.split(".")
+    if len(parts) != 3:
+        raise ValueError("Invalid token format")
+
+    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

hatchet_sdk/v0/utils/aio_utils.py

@@ -0,0 +1,137 @@
+import asyncio
+import inspect
+from concurrent.futures import Executor
+from functools import partial, wraps
+from threading import Thread
+from typing import Any
+
+
+## TODO: Stricter typing here
+def sync_to_async(func: Any) -> Any:
+    """
+    A decorator to run a synchronous function or coroutine in an asynchronous context with added
+    asyncio loop safety.
+
+    This decorator allows you to safely call synchronous functions or coroutines from an
+    asynchronous function by running them in an executor.
+
+    Args:
+        func (callable): The synchronous function or coroutine to be run asynchronously.
+
+    Returns:
+        callable: An asynchronous wrapper function that runs the given function in an executor.
+
+    Example:
+        @sync_to_async
+        def sync_function(x, y):
+            return x + y
+
+        @sync_to_async
+        async def async_function(x, y):
+            return x + y
+
+
+        def undecorated_function(x, y):
+            return x + y
+
+        async def main():
+            result1 = await sync_function(1, 2)
+            result2 = await async_function(3, 4)
+            result3 = await sync_to_async(undecorated_function)(5, 6)
+            print(result1, result2, result3)
+
+        asyncio.run(main())
+    """
+
+    ## TODO: Stricter typing here
+    @wraps(func)
+    async def run(
+        *args: Any,
+        loop: asyncio.AbstractEventLoop | None = None,
+        executor: Executor | None = None,
+        **kwargs: Any
+    ) -> Any:
+        """
+        The asynchronous wrapper function that runs the given function in an executor.
+
+        Args:
+            *args: Positional arguments to pass to the function.
+            loop (asyncio.AbstractEventLoop, optional): The event loop to use. If None, the current running loop is used.
+            executor (concurrent.futures.Executor, optional): The executor to use. If None, the default executor is used.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            The result of the function call.
+        """
+        if loop is None:
+            loop = asyncio.get_running_loop()
+
+        if inspect.iscoroutinefunction(func):
+            # Wrap the coroutine to run it in an executor
+            async def wrapper() -> Any:
+                return await func(*args, **kwargs)
+
+            pfunc = partial(asyncio.run, wrapper())
+            return await loop.run_in_executor(executor, pfunc)
+        else:
+            # Run the synchronous function in an executor
+            pfunc = partial(func, *args, **kwargs)
+            return await loop.run_in_executor(executor, pfunc)
+
+    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.
+
+    Returns:
+        asyncio.AbstractEventLoop: The active event loop, or None if there is no active
+        event loop in the current thread.
+    """
+    try:
+        return asyncio.get_event_loop()
+    except RuntimeError as e:
+        if str(e).startswith("There is no current event loop in thread"):
+            return None
+        else:
+            raise e

hatchet_sdk/v0/utils/backoff.py

@@ -0,0 +1,9 @@
+import asyncio
+import random
+
+
+async def exp_backoff_sleep(attempt: int, max_sleep_time: float = 5) -> None:
+    base_time = 0.1  # starting sleep time in seconds (100 milliseconds)
+    jitter = random.uniform(0, base_time)  # add random jitter
+    sleep_time = min(base_time * (2**attempt) + jitter, max_sleep_time)
+    await asyncio.sleep(sleep_time)