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

hatchet_sdk/v0/clients/admin.py

@@ -0,0 +1,541 @@
+import json
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, TypedDict, TypeVar, Union
+
+import grpc
+from google.protobuf import timestamp_pb2
+
+from hatchet_sdk.v0.clients.rest.models.workflow_run import WorkflowRun
+from hatchet_sdk.v0.clients.rest.tenacity_utils import tenacity_retry
+from hatchet_sdk.v0.clients.run_event_listener import new_listener
+from hatchet_sdk.v0.clients.workflow_listener import PooledWorkflowRunListener
+from hatchet_sdk.v0.connection import new_conn
+from hatchet_sdk.v0.contracts.workflows_pb2 import (
+    BulkTriggerWorkflowRequest,
+    BulkTriggerWorkflowResponse,
+    CreateWorkflowVersionOpts,
+    PutRateLimitRequest,
+    PutWorkflowRequest,
+    RateLimitDuration,
+    ScheduleWorkflowRequest,
+    TriggerWorkflowRequest,
+    TriggerWorkflowResponse,
+    WorkflowVersion,
+)
+from hatchet_sdk.v0.contracts.workflows_pb2_grpc import WorkflowServiceStub
+from hatchet_sdk.v0.workflow_run import RunRef, WorkflowRunRef
+
+from ..loader import ClientConfig
+from ..metadata import get_metadata
+from ..workflow import WorkflowMeta
+
+
+def new_admin(config: ClientConfig):
+    return AdminClient(config)
+
+
+class ScheduleTriggerWorkflowOptions(TypedDict, total=False):
+    parent_id: Optional[str]
+    parent_step_run_id: Optional[str]
+    child_index: Optional[int]
+    child_key: Optional[str]
+    namespace: Optional[str]
+
+
+class ChildTriggerWorkflowOptions(TypedDict, total=False):
+    additional_metadata: Dict[str, str] | None = None
+    sticky: bool | None = None
+
+
+class ChildWorkflowRunDict(TypedDict, total=False):
+    workflow_name: str
+    input: Any
+    options: ChildTriggerWorkflowOptions
+    key: str | None = None
+
+
+class TriggerWorkflowOptions(ScheduleTriggerWorkflowOptions, total=False):
+    additional_metadata: Dict[str, str] | None = None
+    desired_worker_id: str | None = None
+    namespace: str | None = None
+
+
+class WorkflowRunDict(TypedDict, total=False):
+    workflow_name: str
+    input: Any
+    options: TriggerWorkflowOptions | None
+
+
+class DedupeViolationErr(Exception):
+    """Raised by the Hatchet library to indicate that a workflow has already been run with this deduplication value."""
+
+    pass
+
+
+class AdminClientBase:
+    pooled_workflow_listener: PooledWorkflowRunListener | None = None
+
+    def _prepare_workflow_request(
+        self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
+    ):
+        try:
+            payload_data = json.dumps(input)
+
+            try:
+                meta = (
+                    None
+                    if options is None or "additional_metadata" not in options
+                    else options["additional_metadata"]
+                )
+                if meta is not None:
+                    options = {
+                        **options,
+                        "additional_metadata": json.dumps(meta).encode("utf-8"),
+                    }
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Error encoding payload: {e}")
+
+            return TriggerWorkflowRequest(
+                name=workflow_name, input=payload_data, **(options or {})
+            )
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Error encoding payload: {e}")
+
+    def _prepare_put_workflow_request(
+        self,
+        name: str,
+        workflow: CreateWorkflowVersionOpts | WorkflowMeta,
+        overrides: CreateWorkflowVersionOpts | None = None,
+    ):
+        try:
+            opts: CreateWorkflowVersionOpts
+
+            if isinstance(workflow, CreateWorkflowVersionOpts):
+                opts = workflow
+            else:
+                opts = workflow.get_create_opts(self.client.config.namespace)
+
+            if overrides is not None:
+                opts.MergeFrom(overrides)
+
+            opts.name = name
+
+            return PutWorkflowRequest(
+                opts=opts,
+            )
+        except grpc.RpcError as e:
+            raise ValueError(f"Could not put workflow: {e}")
+
+    def _prepare_schedule_workflow_request(
+        self,
+        name: str,
+        schedules: List[Union[datetime, timestamp_pb2.Timestamp]],
+        input={},
+        options: ScheduleTriggerWorkflowOptions = None,
+    ):
+        timestamp_schedules = []
+        for schedule in schedules:
+            if isinstance(schedule, datetime):
+                t = schedule.timestamp()
+                seconds = int(t)
+                nanos = int(t % 1 * 1e9)
+                timestamp = timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
+                timestamp_schedules.append(timestamp)
+            elif isinstance(schedule, timestamp_pb2.Timestamp):
+                timestamp_schedules.append(schedule)
+            else:
+                raise ValueError(
+                    "Invalid schedule type. Must be datetime or timestamp_pb2.Timestamp."
+                )
+
+        return ScheduleWorkflowRequest(
+            name=name,
+            schedules=timestamp_schedules,
+            input=json.dumps(input),
+            **(options or {}),
+        )
+
+
+T = TypeVar("T")
+
+
+class AdminClientAioImpl(AdminClientBase):
+    def __init__(self, config: ClientConfig):
+        aio_conn = new_conn(config, True)
+        self.config = config
+        self.aio_client = WorkflowServiceStub(aio_conn)
+        self.token = config.token
+        self.listener_client = new_listener(config)
+        self.namespace = config.namespace
+
+    async def run(
+        self,
+        function: Union[str, Callable[[Any], T]],
+        input: any,
+        options: TriggerWorkflowOptions = None,
+    ) -> "RunRef[T]":
+        workflow_name = function
+
+        if not isinstance(function, str):
+            workflow_name = function.function_name
+
+        wrr = await self.run_workflow(workflow_name, input, options)
+
+        return RunRef[T](
+            wrr.workflow_run_id, wrr.workflow_listener, wrr.workflow_run_event_listener
+        )
+
+    ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
+    @tenacity_retry
+    async def run_workflow(
+        self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
+    ) -> WorkflowRunRef:
+        try:
+            if not self.pooled_workflow_listener:
+                self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
+            namespace = self.namespace
+
+            if (
+                options is not None
+                and "namespace" in options
+                and options["namespace"] is not None
+            ):
+                namespace = options.pop("namespace")
+
+            if namespace != "" and not workflow_name.startswith(self.namespace):
+                workflow_name = f"{namespace}{workflow_name}"
+
+            request = self._prepare_workflow_request(workflow_name, input, options)
+
+            resp: TriggerWorkflowResponse = await self.aio_client.TriggerWorkflow(
+                request,
+                metadata=get_metadata(self.token),
+            )
+
+            return WorkflowRunRef(
+                workflow_run_id=resp.workflow_run_id,
+                workflow_listener=self.pooled_workflow_listener,
+                workflow_run_event_listener=self.listener_client,
+            )
+        except (grpc.RpcError, grpc.aio.AioRpcError) as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise DedupeViolationErr(e.details())
+
+            raise e
+
+    ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
+    @tenacity_retry
+    async def run_workflows(
+        self,
+        workflows: list[WorkflowRunDict],
+        options: TriggerWorkflowOptions | None = None,
+    ) -> List[WorkflowRunRef]:
+        if len(workflows) == 0:
+            raise ValueError("No workflows to run")
+
+        if not self.pooled_workflow_listener:
+            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
+        namespace = self.namespace
+
+        if (
+            options is not None
+            and "namespace" in options
+            and options["namespace"] is not None
+        ):
+            namespace = options["namespace"]
+            del options["namespace"]
+
+        workflow_run_requests: TriggerWorkflowRequest = []
+
+        for workflow in workflows:
+            workflow_name = workflow["workflow_name"]
+            input_data = workflow["input"]
+            options = workflow["options"]
+
+            if namespace != "" and not workflow_name.startswith(self.namespace):
+                workflow_name = f"{namespace}{workflow_name}"
+
+            # Prepare and trigger workflow for each workflow name and input
+            request = self._prepare_workflow_request(workflow_name, input_data, options)
+            workflow_run_requests.append(request)
+
+        request = BulkTriggerWorkflowRequest(workflows=workflow_run_requests)
+
+        resp: BulkTriggerWorkflowResponse = await self.aio_client.BulkTriggerWorkflow(
+            request,
+            metadata=get_metadata(self.token),
+        )
+
+        return [
+            WorkflowRunRef(
+                workflow_run_id=workflow_run_id,
+                workflow_listener=self.pooled_workflow_listener,
+                workflow_run_event_listener=self.listener_client,
+            )
+            for workflow_run_id in resp.workflow_run_ids
+        ]
+
+    @tenacity_retry
+    async def put_workflow(
+        self,
+        name: str,
+        workflow: CreateWorkflowVersionOpts | WorkflowMeta,
+        overrides: CreateWorkflowVersionOpts | None = None,
+    ) -> WorkflowVersion:
+        opts = self._prepare_put_workflow_request(name, workflow, overrides)
+
+        return await self.aio_client.PutWorkflow(
+            opts,
+            metadata=get_metadata(self.token),
+        )
+
+    @tenacity_retry
+    async def put_rate_limit(
+        self,
+        key: str,
+        limit: int,
+        duration: RateLimitDuration = RateLimitDuration.SECOND,
+    ):
+        await self.aio_client.PutRateLimit(
+            PutRateLimitRequest(
+                key=key,
+                limit=limit,
+                duration=duration,
+            ),
+            metadata=get_metadata(self.token),
+        )
+
+    @tenacity_retry
+    async def schedule_workflow(
+        self,
+        name: str,
+        schedules: List[Union[datetime, timestamp_pb2.Timestamp]],
+        input={},
+        options: ScheduleTriggerWorkflowOptions = None,
+    ) -> WorkflowVersion:
+        try:
+            namespace = self.namespace
+
+            if (
+                options is not None
+                and "namespace" in options
+                and options["namespace"] is not None
+            ):
+                namespace = options["namespace"]
+                del options["namespace"]
+
+            if namespace != "" and not name.startswith(self.namespace):
+                name = f"{namespace}{name}"
+
+            request = self._prepare_schedule_workflow_request(
+                name, schedules, input, options
+            )
+
+            return await self.aio_client.ScheduleWorkflow(
+                request,
+                metadata=get_metadata(self.token),
+            )
+        except (grpc.aio.AioRpcError, grpc.RpcError) as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise DedupeViolationErr(e.details())
+
+            raise e
+
+
+class AdminClient(AdminClientBase):
+    def __init__(self, config: ClientConfig):
+        conn = new_conn(config)
+        self.config = config
+        self.client = WorkflowServiceStub(conn)
+        self.aio = AdminClientAioImpl(config)
+        self.token = config.token
+        self.listener_client = new_listener(config)
+        self.namespace = config.namespace
+
+    @tenacity_retry
+    def put_workflow(
+        self,
+        name: str,
+        workflow: CreateWorkflowVersionOpts | WorkflowMeta,
+        overrides: CreateWorkflowVersionOpts | None = None,
+    ) -> WorkflowVersion:
+        opts = self._prepare_put_workflow_request(name, workflow, overrides)
+
+        resp: WorkflowVersion = self.client.PutWorkflow(
+            opts,
+            metadata=get_metadata(self.token),
+        )
+
+        return resp
+
+    @tenacity_retry
+    def put_rate_limit(
+        self,
+        key: str,
+        limit: int,
+        duration: Union[RateLimitDuration.Value, str] = RateLimitDuration.SECOND,
+    ):
+        self.client.PutRateLimit(
+            PutRateLimitRequest(
+                key=key,
+                limit=limit,
+                duration=duration,
+            ),
+            metadata=get_metadata(self.token),
+        )
+
+    @tenacity_retry
+    def schedule_workflow(
+        self,
+        name: str,
+        schedules: List[Union[datetime, timestamp_pb2.Timestamp]],
+        input={},
+        options: ScheduleTriggerWorkflowOptions = None,
+    ) -> WorkflowVersion:
+        try:
+            namespace = self.namespace
+
+            if (
+                options is not None
+                and "namespace" in options
+                and options["namespace"] is not None
+            ):
+                namespace = options["namespace"]
+                del options["namespace"]
+
+            if namespace != "" and not name.startswith(self.namespace):
+                name = f"{namespace}{name}"
+
+            request = self._prepare_schedule_workflow_request(
+                name, schedules, input, options
+            )
+
+            return self.client.ScheduleWorkflow(
+                request,
+                metadata=get_metadata(self.token),
+            )
+        except (grpc.RpcError, grpc.aio.AioRpcError) as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise DedupeViolationErr(e.details())
+
+            raise e
+
+    ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
+    @tenacity_retry
+    def run_workflow(
+        self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
+    ) -> WorkflowRunRef:
+        try:
+            if not self.pooled_workflow_listener:
+                self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
+            namespace = self.namespace
+
+            ## TODO: Factor this out - it's repeated a lot of places
+            if (
+                options is not None
+                and "namespace" in options
+                and options["namespace"] is not None
+            ):
+                namespace = options.pop("namespace")
+
+            if namespace != "" and not workflow_name.startswith(self.namespace):
+                workflow_name = f"{namespace}{workflow_name}"
+
+            request = self._prepare_workflow_request(workflow_name, input, options)
+
+            resp: TriggerWorkflowResponse = self.client.TriggerWorkflow(
+                request,
+                metadata=get_metadata(self.token),
+            )
+
+            return WorkflowRunRef(
+                workflow_run_id=resp.workflow_run_id,
+                workflow_listener=self.pooled_workflow_listener,
+                workflow_run_event_listener=self.listener_client,
+            )
+        except (grpc.RpcError, grpc.aio.AioRpcError) as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise DedupeViolationErr(e.details())
+
+            raise e
+
+    ## IMPORTANT: Keep this method's signature in sync with the wrapper in the OTel instrumentor
+    @tenacity_retry
+    def run_workflows(
+        self, workflows: List[WorkflowRunDict], options: TriggerWorkflowOptions = None
+    ) -> list[WorkflowRunRef]:
+        workflow_run_requests: TriggerWorkflowRequest = []
+        if not self.pooled_workflow_listener:
+            self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
+        for workflow in workflows:
+            workflow_name = workflow["workflow_name"]
+            input_data = workflow["input"]
+            options = workflow["options"]
+
+            namespace = self.namespace
+
+            if (
+                options is not None
+                and "namespace" in options
+                and options["namespace"] is not None
+            ):
+                namespace = options["namespace"]
+                del options["namespace"]
+
+            if namespace != "" and not workflow_name.startswith(self.namespace):
+                workflow_name = f"{namespace}{workflow_name}"
+
+            # Prepare and trigger workflow for each workflow name and input
+            request = self._prepare_workflow_request(workflow_name, input_data, options)
+
+            workflow_run_requests.append(request)
+
+            request = BulkTriggerWorkflowRequest(workflows=workflow_run_requests)
+
+        resp: BulkTriggerWorkflowResponse = self.client.BulkTriggerWorkflow(
+            request,
+            metadata=get_metadata(self.token),
+        )
+
+        return [
+            WorkflowRunRef(
+                workflow_run_id=workflow_run_id,
+                workflow_listener=self.pooled_workflow_listener,
+                workflow_run_event_listener=self.listener_client,
+            )
+            for workflow_run_id in resp.workflow_run_ids
+        ]
+
+    def run(
+        self,
+        function: Union[str, Callable[[Any], T]],
+        input: any,
+        options: TriggerWorkflowOptions = None,
+    ) -> "RunRef[T]":
+        workflow_name = function
+
+        if not isinstance(function, str):
+            workflow_name = function.function_name
+
+        wrr = self.run_workflow(workflow_name, input, options)
+
+        return RunRef[T](
+            wrr.workflow_run_id, wrr.workflow_listener, wrr.workflow_run_event_listener
+        )
+
+    def get_workflow_run(self, workflow_run_id: str) -> WorkflowRunRef:
+        try:
+            if not self.pooled_workflow_listener:
+                self.pooled_workflow_listener = PooledWorkflowRunListener(self.config)
+
+            return WorkflowRunRef(
+                workflow_run_id=workflow_run_id,
+                workflow_listener=self.pooled_workflow_listener,
+                workflow_run_event_listener=self.listener_client,
+            )
+        except grpc.RpcError as e:
+            raise ValueError(f"Could not get workflow run: {e}")