path: root/.venv/lib/python3.12/site-packages/hatchet_sdk/context/context.py

import inspect
import json
import traceback
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Any, Generic, Type, TypeVar, cast, overload
from warnings import warn

from pydantic import BaseModel, StrictStr

from hatchet_sdk.clients.events import EventClient
from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
from hatchet_sdk.clients.rest_client import RestApi
from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
from hatchet_sdk.context.worker_context import WorkerContext
from hatchet_sdk.contracts.dispatcher_pb2 import OverridesData
from hatchet_sdk.contracts.workflows_pb2 import (
    BulkTriggerWorkflowRequest,
    TriggerWorkflowRequest,
)
from hatchet_sdk.utils.types import WorkflowValidator
from hatchet_sdk.utils.typing import is_basemodel_subclass
from hatchet_sdk.workflow_run import WorkflowRunRef

from ..clients.admin import (
    AdminClient,
    ChildTriggerWorkflowOptions,
    ChildWorkflowRunDict,
    TriggerWorkflowOptions,
    WorkflowRunDict,
)
from ..clients.dispatcher.dispatcher import (  # type: ignore[attr-defined]
    Action,
    DispatcherClient,
)
from ..logger import logger

DEFAULT_WORKFLOW_POLLING_INTERVAL = 5  # Seconds

T = TypeVar("T", bound=BaseModel)


def get_caller_file_path() -> str:
    caller_frame = inspect.stack()[2]

    return caller_frame.filename


class BaseContext:

    action: Action
    spawn_index: int

    def _prepare_workflow_options(
        self,
        key: str | None = None,
        options: ChildTriggerWorkflowOptions | None = None,
        worker_id: str | None = None,
    ) -> TriggerWorkflowOptions:
        workflow_run_id = self.action.workflow_run_id
        step_run_id = self.action.step_run_id

        desired_worker_id = None
        if options is not None and "sticky" in options and options["sticky"] == True:
            desired_worker_id = worker_id

        meta = None
        if options is not None and "additional_metadata" in options:
            meta = options["additional_metadata"]

        ## TODO: Pydantic here to simplify this
        trigger_options: TriggerWorkflowOptions = {
            "parent_id": workflow_run_id,
            "parent_step_run_id": step_run_id,
            "child_key": key,
            "child_index": self.spawn_index,
            "additional_metadata": meta,
            "desired_worker_id": desired_worker_id,
        }

        self.spawn_index += 1
        return trigger_options


class ContextAioImpl(BaseContext):
    def __init__(
        self,
        action: Action,
        dispatcher_client: DispatcherClient,
        admin_client: AdminClient,
        event_client: EventClient,
        rest_client: RestApi,
        workflow_listener: PooledWorkflowRunListener,
        workflow_run_event_listener: RunEventListenerClient,
        worker: WorkerContext,
        namespace: str = "",
    ):
        self.action = action
        self.dispatcher_client = dispatcher_client
        self.admin_client = admin_client
        self.event_client = event_client
        self.rest_client = rest_client
        self.workflow_listener = workflow_listener
        self.workflow_run_event_listener = workflow_run_event_listener
        self.namespace = namespace
        self.spawn_index = -1
        self.worker = worker

    @tenacity_retry
    async def spawn_workflow(
        self,
        workflow_name: str,
        input: dict[str, Any] = {},
        key: str | None = None,
        options: ChildTriggerWorkflowOptions | None = None,
    ) -> WorkflowRunRef:
        worker_id = self.worker.id()
        # if (
        #     options is not None
        #     and "sticky" in options
        #     and options["sticky"] == True
        #     and not self.worker.has_workflow(workflow_name)
        # ):
        #     raise Exception(
        #         f"cannot run with sticky: workflow {workflow_name} is not registered on the worker"
        #     )

        trigger_options = self._prepare_workflow_options(key, options, worker_id)

        return await self.admin_client.aio.run_workflow(
            workflow_name, input, trigger_options
        )

    @tenacity_retry
    async def spawn_workflows(
        self, child_workflow_runs: list[ChildWorkflowRunDict]
    ) -> list[WorkflowRunRef]:

        if len(child_workflow_runs) == 0:
            raise Exception("no child workflows to spawn")

        worker_id = self.worker.id()

        bulk_trigger_workflow_runs: list[WorkflowRunDict] = []
        for child_workflow_run in child_workflow_runs:
            workflow_name = child_workflow_run["workflow_name"]
            input = child_workflow_run["input"]

            key = child_workflow_run.get("key")
            options = child_workflow_run.get("options", {})

            trigger_options = self._prepare_workflow_options(key, options, worker_id)

            bulk_trigger_workflow_runs.append(
                WorkflowRunDict(
                    workflow_name=workflow_name, input=input, options=trigger_options
                )
            )

        return await self.admin_client.aio.run_workflows(bulk_trigger_workflow_runs)


class Context(BaseContext):
    spawn_index = -1

    worker: WorkerContext

    def __init__(
        self,
        action: Action,
        dispatcher_client: DispatcherClient,
        admin_client: AdminClient,
        event_client: EventClient,
        rest_client: RestApi,
        workflow_listener: PooledWorkflowRunListener,
        workflow_run_event_listener: RunEventListenerClient,
        worker: WorkerContext,
        namespace: str = "",
        validator_registry: dict[str, WorkflowValidator] = {},
    ):
        self.worker = worker
        self.validator_registry = validator_registry

        self.aio = ContextAioImpl(
            action,
            dispatcher_client,
            admin_client,
            event_client,
            rest_client,
            workflow_listener,
            workflow_run_event_listener,
            worker,
            namespace,
        )

        # Check the type of action.action_payload before attempting to load it as JSON
        if isinstance(action.action_payload, (str, bytes, bytearray)):
            try:
                self.data = cast(dict[str, Any], json.loads(action.action_payload))
            except Exception as e:
                logger.error(f"Error parsing action payload: {e}")
                # Assign an empty dictionary if parsing fails
                self.data: dict[str, Any] = {}  # type: ignore[no-redef]
        else:
            # Directly assign the payload to self.data if it's already a dict
            self.data = (
                action.action_payload if isinstance(action.action_payload, dict) else {}
            )

        self.action = action

        # FIXME: stepRunId is a legacy field, we should remove it
        self.stepRunId = action.step_run_id

        self.step_run_id = action.step_run_id
        self.exit_flag = False
        self.dispatcher_client = dispatcher_client
        self.admin_client = admin_client
        self.event_client = event_client
        self.rest_client = rest_client
        self.workflow_listener = workflow_listener
        self.workflow_run_event_listener = workflow_run_event_listener
        self.namespace = namespace

        # FIXME: this limits the number of concurrent log requests to 1, which means we can do about
        # 100 log lines per second but this depends on network.
        self.logger_thread_pool = ThreadPoolExecutor(max_workers=1)
        self.stream_event_thread_pool = ThreadPoolExecutor(max_workers=1)

        # store each key in the overrides field in a lookup table
        # overrides_data is a dictionary of key-value pairs
        self.overrides_data = self.data.get("overrides", {})

        if action.get_group_key_run_id != "":
            self.input = self.data
        else:
            self.input = self.data.get("input", {})

    def step_output(self, step: str) -> dict[str, Any] | BaseModel:
        workflow_validator = next(
            (v for k, v in self.validator_registry.items() if k.split(":")[-1] == step),
            None,
        )

        try:
            parent_step_data = cast(dict[str, Any], self.data["parents"][step])
        except KeyError:
            raise ValueError(f"Step output for '{step}' not found")

        if workflow_validator and (v := workflow_validator.step_output):
            return v.model_validate(parent_step_data)

        return parent_step_data

    def triggered_by_event(self) -> bool:
        return cast(str, self.data.get("triggered_by", "")) == "event"

    def workflow_input(self) -> dict[str, Any] | T:
        if (r := self.validator_registry.get(self.action.action_id)) and (
            i := r.workflow_input
        ):
            return cast(
                T,
                i.model_validate(self.input),
            )

        return self.input

    def workflow_run_id(self) -> str:
        return self.action.workflow_run_id

    def cancel(self) -> None:
        logger.debug("cancelling step...")
        self.exit_flag = True

    # done returns true if the context has been cancelled
    def done(self) -> bool:
        return self.exit_flag

    def playground(self, name: str, default: str | None = None) -> str | None:
        # if the key exists in the overrides_data field, return the value
        if name in self.overrides_data:
            warn(
                "Use of `overrides_data` is deprecated.",
                DeprecationWarning,
                stacklevel=1,
            )
            return str(self.overrides_data[name])

        caller_file = get_caller_file_path()

        self.dispatcher_client.put_overrides_data(
            OverridesData(
                stepRunId=self.stepRunId,
                path=name,
                value=json.dumps(default),
                callerFilename=caller_file,
            )
        )

        return default

    def _log(self, line: str) -> tuple[bool, Exception | None]:
        try:
            self.event_client.log(message=line, step_run_id=self.stepRunId)
            return True, None
        except Exception as e:
            # we don't want to raise an exception here, as it will kill the log thread
            return False, e

    def log(self, line: Any, raise_on_error: bool = False) -> None:
        if self.stepRunId == "":
            return

        if not isinstance(line, str):
            try:
                line = json.dumps(line)
            except Exception:
                line = str(line)

        future = self.logger_thread_pool.submit(self._log, line)

        def handle_result(future: Future[tuple[bool, Exception | None]]) -> None:
            success, exception = future.result()
            if not success and exception:
                if raise_on_error:
                    raise exception
                else:
                    thread_trace = "".join(
                        traceback.format_exception(
                            type(exception), exception, exception.__traceback__
                        )
                    )
                    call_site_trace = "".join(traceback.format_stack())
                    logger.error(
                        f"Error in log thread: {exception}\n{thread_trace}\nCalled from:\n{call_site_trace}"
                    )

        future.add_done_callback(handle_result)

    def release_slot(self) -> None:
        return self.dispatcher_client.release_slot(self.stepRunId)

    def _put_stream(self, data: str | bytes) -> None:
        try:
            self.event_client.stream(data=data, step_run_id=self.stepRunId)
        except Exception as e:
            logger.error(f"Error putting stream event: {e}")

    def put_stream(self, data: str | bytes) -> None:
        if self.stepRunId == "":
            return

        self.stream_event_thread_pool.submit(self._put_stream, data)

    def refresh_timeout(self, increment_by: str) -> None:
        try:
            return self.dispatcher_client.refresh_timeout(
                step_run_id=self.stepRunId, increment_by=increment_by
            )
        except Exception as e:
            logger.error(f"Error refreshing timeout: {e}")

    def retry_count(self) -> int:
        return self.action.retry_count

    def additional_metadata(self) -> dict[str, Any] | None:
        return self.action.additional_metadata

    def child_index(self) -> int | None:
        return self.action.child_workflow_index

    def child_key(self) -> str | None:
        return self.action.child_workflow_key

    def parent_workflow_run_id(self) -> str | None:
        return self.action.parent_workflow_run_id

    def step_run_errors(self) -> dict[str, str]:
        errors = cast(dict[str, str], self.data.get("step_run_errors", {}))

        if not errors:
            logger.error(
                "No step run errors found. `context.step_run_errors` is intended to be run in an on-failure step, and will only work on engine versions more recent than v0.53.10"
            )

        return errors

    def fetch_run_failures(self) -> list[dict[str, StrictStr]]:
        data = self.rest_client.workflow_run_get(self.action.workflow_run_id)
        other_job_runs = [
            run for run in (data.job_runs or []) if run.job_id != self.action.job_id
        ]
        # TODO: Parse Step Runs using a Pydantic Model rather than a hand crafted dictionary
        return [
            {
                "step_id": step_run.step_id,
                "step_run_action_name": step_run.step.action,
                "error": step_run.error,
            }
            for job_run in other_job_runs
            if job_run.step_runs
            for step_run in job_run.step_runs
            if step_run.error and step_run.step
        ]

    @tenacity_retry
    def spawn_workflow(
        self,
        workflow_name: str,
        input: dict[str, Any] = {},
        key: str | None = None,
        options: ChildTriggerWorkflowOptions | None = None,
    ) -> WorkflowRunRef:
        worker_id = self.worker.id()
        trigger_options = self._prepare_workflow_options(key, options, worker_id)

        return self.admin_client.run_workflow(workflow_name, input, trigger_options)

    @tenacity_retry
    def spawn_workflows(
        self, child_workflow_runs: list[ChildWorkflowRunDict]
    ) -> list[WorkflowRunRef]:

        if len(child_workflow_runs) == 0:
            raise Exception("no child workflows to spawn")

        worker_id = self.worker.id()

        bulk_trigger_workflow_runs: list[WorkflowRunDict] = []
        for child_workflow_run in child_workflow_runs:
            workflow_name = child_workflow_run["workflow_name"]
            input = child_workflow_run["input"]

            key = child_workflow_run.get("key")
            options = child_workflow_run.get("options", {})

            trigger_options = self._prepare_workflow_options(key, options, worker_id)

            bulk_trigger_workflow_runs.append(
                WorkflowRunDict(
                    workflow_name=workflow_name, input=input, options=trigger_options
                )
            )

        return self.admin_client.run_workflows(bulk_trigger_workflow_runs)