hatchet-sdk 0.47.0__py3-none-any.whl → 1.0.0__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/semver.py

@@ -0,0 +1,30 @@
+def bump_minor_version(version: str) -> str:
+    """
+    Bumps the minor version of a semantic version string. NOTE this doesn't follow full semver,
+    missing the build metadata and pre-release version.
+
+    :param version: A semantic version string in the format major.minor.patch
+    :return: A string with the minor version bumped and patch version reset to 0
+    :raises ValueError: If the input is not a valid semantic version string
+    """
+    # if it starts with a v, remove it
+    had_v = False
+    if version.startswith("v"):
+        version = version[1:]
+        had_v = True
+
+    parts = version.split(".")
+    if len(parts) != 3:
+        raise ValueError(f"Invalid semantic version: {version}")
+
+    try:
+        major, minor, _ = map(int, parts)
+    except ValueError:
+        raise ValueError(f"Invalid semantic version: {version}")
+
+    new_minor = minor + 1
+    new_version = f"{major}.{new_minor}.0"
+
+    if had_v:
+        new_version = "v" + new_version
+    return new_version

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