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

hatchet_sdk/workflow.py

@@ -1,527 +0,0 @@
-import asyncio
-from dataclasses import dataclass, field
-from datetime import datetime
-from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Awaitable,
-    Callable,
-    Generic,
-    ParamSpec,
-    Type,
-    TypeGuard,
-    TypeVar,
-    Union,
-    cast,
-)
-
-from google.protobuf import timestamp_pb2
-from pydantic import BaseModel, ConfigDict
-
-from hatchet_sdk.clients.admin import (
-    ChildTriggerWorkflowOptions,
-    ChildWorkflowRunDict,
-    ScheduleTriggerWorkflowOptions,
-)
-from hatchet_sdk.context.context import Context
-from hatchet_sdk.contracts.workflows_pb2 import (
-    ConcurrencyLimitStrategy as ConcurrencyLimitStrategyProto,
-)
-from hatchet_sdk.contracts.workflows_pb2 import (
-    CreateStepRateLimit,
-    CreateWorkflowJobOpts,
-    CreateWorkflowStepOpts,
-    CreateWorkflowVersionOpts,
-    DesiredWorkerLabels,
-)
-from hatchet_sdk.contracts.workflows_pb2 import StickyStrategy as StickyStrategyProto
-from hatchet_sdk.contracts.workflows_pb2 import (
-    WorkflowConcurrencyOpts,
-    WorkflowKind,
-    WorkflowVersion,
-)
-from hatchet_sdk.labels import DesiredWorkerLabel
-from hatchet_sdk.logger import logger
-from hatchet_sdk.rate_limit import RateLimit
-from hatchet_sdk.utils.proto_enums import convert_python_enum_to_proto, maybe_int_to_str
-from hatchet_sdk.workflow_run import WorkflowRunRef
-
-if TYPE_CHECKING:
-    from hatchet_sdk import Hatchet
-
-R = TypeVar("R")
-P = ParamSpec("P")
-
-
-class EmptyModel(BaseModel):
-    model_config = ConfigDict(extra="allow")
-
-
-class StickyStrategy(str, Enum):
-    SOFT = "SOFT"
-    HARD = "HARD"
-
-
-class ConcurrencyLimitStrategy(str, Enum):
-    CANCEL_IN_PROGRESS = "CANCEL_IN_PROGRESS"
-    DROP_NEWEST = "DROP_NEWEST"
-    QUEUE_NEWEST = "QUEUE_NEWEST"
-    GROUP_ROUND_ROBIN = "GROUP_ROUND_ROBIN"
-    CANCEL_NEWEST = "CANCEL_NEWEST"
-
-
-class ConcurrencyExpression(BaseModel):
-    """
-    Defines concurrency limits for a workflow using a CEL expression.
-
-    Args:
-        expression (str): CEL expression to determine concurrency grouping. (i.e. "input.user_id")
-        max_runs (int): Maximum number of concurrent workflow runs.
-        limit_strategy (ConcurrencyLimitStrategy): Strategy for handling limit violations.
-
-    Example:
-        ConcurrencyExpression("input.user_id", 5, ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS)
-    """
-
-    expression: str
-    max_runs: int
-    limit_strategy: ConcurrencyLimitStrategy
-
-
-TWorkflowInput = TypeVar("TWorkflowInput", bound=BaseModel)
-
-
-class WorkflowConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-
-    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[BaseModel] = EmptyModel
-
-
-class StepType(str, Enum):
-    DEFAULT = "default"
-    CONCURRENCY = "concurrency"
-    ON_FAILURE = "on_failure"
-
-
-AsyncFunc = Callable[[Any, Context], Awaitable[R]]
-SyncFunc = Callable[[Any, Context], R]
-StepFunc = Union[AsyncFunc[R], SyncFunc[R]]
-
-
-def is_async_fn(fn: StepFunc[R]) -> TypeGuard[AsyncFunc[R]]:
-    return asyncio.iscoroutinefunction(fn)
-
-
-def is_sync_fn(fn: StepFunc[R]) -> TypeGuard[SyncFunc[R]]:
-    return not asyncio.iscoroutinefunction(fn)
-
-
-class Step(Generic[R]):
-    def __init__(
-        self,
-        fn: Callable[[Any, Context], R] | Callable[[Any, Context], Awaitable[R]],
-        type: StepType,
-        name: str = "",
-        timeout: str = "60m",
-        parents: list[str] = [],
-        retries: int = 0,
-        rate_limits: list[CreateStepRateLimit] = [],
-        desired_worker_labels: dict[str, DesiredWorkerLabels] = {},
-        backoff_factor: float | None = None,
-        backoff_max_seconds: int | None = None,
-        concurrency__max_runs: int | None = None,
-        concurrency__limit_strategy: ConcurrencyLimitStrategy | None = None,
-    ) -> None:
-        self.fn = fn
-        self.is_async_function = is_async_fn(fn)
-        self.workflow: Union["BaseWorkflow", None] = None
-
-        self.type = type
-        self.timeout = timeout
-        self.name = name
-        self.parents = parents
-        self.retries = retries
-        self.rate_limits = rate_limits
-        self.desired_worker_labels = desired_worker_labels
-        self.backoff_factor = backoff_factor
-        self.backoff_max_seconds = backoff_max_seconds
-        self.concurrency__max_runs = concurrency__max_runs
-        self.concurrency__limit_strategy = concurrency__limit_strategy
-
-    def call(self, ctx: Context) -> R:
-        if not self.is_registered:
-            raise ValueError(
-                "Only steps that have been registered can be called. To register this step, instantiate its corresponding workflow."
-            )
-
-        if self.is_async_function:
-            raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
-
-        sync_fn = self.fn
-        if is_sync_fn(sync_fn):
-            return sync_fn(self.workflow, ctx)
-
-        raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
-
-    async def aio_call(self, ctx: Context) -> R:
-        if not self.is_registered:
-            raise ValueError(
-                "Only steps that have been registered can be called. To register this step, instantiate its corresponding workflow."
-            )
-
-        if not self.is_async_function:
-            raise TypeError(
-                f"{self.name} is not an async function. Use `call` instead."
-            )
-
-        async_fn = self.fn
-
-        if is_async_fn(async_fn):
-            return await async_fn(self.workflow, ctx)
-
-        raise TypeError(f"{self.name} is not an async function. Use `call` instead.")
-
-    @property
-    def is_registered(self) -> bool:
-        return self.workflow is not None
-
-
-class Task(Generic[R, TWorkflowInput]):
-    def __init__(
-        self,
-        fn: Callable[[Context], R],
-        hatchet: "Hatchet",
-        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: Union["Task[R, TWorkflowInput]", None] = None,
-        default_priority: int = 1,
-        input_validator: Type[TWorkflowInput] | None = None,
-        backoff_factor: float | None = None,
-        backoff_max_seconds: int | None = None,
-    ) -> None:
-        def func(_: Any, context: Context) -> R:
-            return fn(context)
-
-        self.hatchet = hatchet
-        self.step: Step[R] = hatchet.step(
-            name=name or fn.__name__,
-            timeout=timeout,
-            retries=retries,
-            rate_limits=rate_limits,
-            desired_worker_labels=desired_worker_labels,
-            backoff_factor=backoff_factor,
-            backoff_max_seconds=backoff_max_seconds,
-        )(func)
-        self.on_failure_step = on_failure
-        self.workflow_config = WorkflowConfig(
-            name=name or fn.__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),
-        )
-
-
-@dataclass
-class SpawnWorkflowInput(Generic[TWorkflowInput]):
-    input: TWorkflowInput
-    key: str | None = None
-    options: ChildTriggerWorkflowOptions = field(
-        default_factory=ChildTriggerWorkflowOptions
-    )
-
-
-class WorkflowDeclaration(Generic[TWorkflowInput]):
-    def __init__(self, config: WorkflowConfig, hatchet: Union["Hatchet", None]):
-        self.config = config
-        self.hatchet = hatchet
-
-    def run(self, input: TWorkflowInput | None = None) -> WorkflowRunRef:
-        if not self.hatchet:
-            raise ValueError("Hatchet client is not initialized.")
-
-        return self.hatchet.admin.run_workflow(
-            workflow_name=self.config.name, input=input.model_dump() if input else {}
-        )
-
-    def get_workflow_input(self, ctx: Context) -> TWorkflowInput:
-        return cast(
-            TWorkflowInput,
-            self.config.input_validator.model_validate(ctx.workflow_input),
-        )
-
-    async def aio_spawn_many(
-        self, ctx: Context, spawn_inputs: list[SpawnWorkflowInput[TWorkflowInput]]
-    ) -> list[WorkflowRunRef]:
-        inputs = [
-            ChildWorkflowRunDict(
-                workflow_name=self.config.name,
-                input=spawn_input.input.model_dump(),
-                key=spawn_input.key,
-                options=spawn_input.options,
-            )
-            for spawn_input in spawn_inputs
-        ]
-        return await ctx.aio_spawn_workflows(inputs)
-
-    async def aio_spawn_one(
-        self,
-        ctx: Context,
-        input: TWorkflowInput,
-        key: str | None = None,
-        options: ChildTriggerWorkflowOptions = ChildTriggerWorkflowOptions(),
-    ) -> WorkflowRunRef:
-        return await ctx.aio_spawn_workflow(
-            workflow_name=self.config.name,
-            input=input.model_dump(),
-            key=key,
-            options=options,
-        )
-
-    def spawn_many(
-        self, ctx: Context, spawn_inputs: list[SpawnWorkflowInput[TWorkflowInput]]
-    ) -> list[WorkflowRunRef]:
-        inputs = [
-            ChildWorkflowRunDict(
-                workflow_name=self.config.name,
-                input=spawn_input.input.model_dump(),
-                key=spawn_input.key,
-                options=spawn_input.options,
-            )
-            for spawn_input in spawn_inputs
-        ]
-
-        return ctx.spawn_workflows(inputs)
-
-    def spawn_one(
-        self,
-        ctx: Context,
-        input: TWorkflowInput,
-        key: str | None = None,
-        options: ChildTriggerWorkflowOptions = ChildTriggerWorkflowOptions(),
-    ) -> WorkflowRunRef:
-        return ctx.spawn_workflow(
-            workflow_name=self.config.name,
-            input=input.model_dump(),
-            key=key,
-            options=options,
-        )
-
-    def schedule(
-        self,
-        schedules: list[datetime | timestamp_pb2.Timestamp],
-        input: TWorkflowInput,
-        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
-    ) -> WorkflowVersion:
-        if not self.hatchet:
-            raise ValueError("Hatchet client is not initialized.")
-
-        return self.hatchet.admin.schedule_workflow(
-            name=self.config.name,
-            schedules=schedules,
-            input=input.model_dump(),
-            options=options,
-        )
-
-    async def aio_schedule(
-        self,
-        schedules: list[datetime | timestamp_pb2.Timestamp],
-        input: TWorkflowInput,
-        options: ScheduleTriggerWorkflowOptions = ScheduleTriggerWorkflowOptions(),
-    ) -> WorkflowVersion:
-        if not self.hatchet:
-            raise ValueError("Hatchet client is not initialized.")
-
-        return await self.hatchet.admin.aio_schedule_workflow(
-            name=self.config.name,
-            schedules=schedules,
-            input=input.model_dump(),
-            options=options,
-        )
-
-
-class BaseWorkflow:
-    """
-    A Hatchet workflow implementation base. This class should be inherited by all workflow implementations.
-
-    Configuration is passed to the workflow implementation via the `config` attribute.
-    """
-
-    config: WorkflowConfig = WorkflowConfig()
-
-    def __init__(self) -> None:
-        self.config.name = self.config.name or str(self.__class__.__name__)
-
-        for step in self.steps:
-            step.workflow = self
-
-    def get_service_name(self, namespace: str) -> str:
-        return f"{namespace}{self.config.name.lower()}"
-
-    def _get_steps_by_type(self, step_type: StepType) -> list[Step[Any]]:
-        return [
-            attr
-            for _, attr in self.__class__.__dict__.items()
-            if isinstance(attr, Step) and attr.type == step_type
-        ]
-
-    @property
-    def on_failure_steps(self) -> list[Step[Any]]:
-        return self._get_steps_by_type(StepType.ON_FAILURE)
-
-    @property
-    def concurrency_actions(self) -> list[Step[Any]]:
-        return self._get_steps_by_type(StepType.CONCURRENCY)
-
-    @property
-    def default_steps(self) -> list[Step[Any]]:
-        return self._get_steps_by_type(StepType.DEFAULT)
-
-    @property
-    def steps(self) -> list[Step[Any]]:
-        return self.default_steps + self.concurrency_actions + self.on_failure_steps
-
-    def create_action_name(self, namespace: str, step: Step[Any]) -> str:
-        return self.get_service_name(namespace) + ":" + step.name
-
-    def get_name(self, namespace: str) -> str:
-        return namespace + self.config.name
-
-    def validate_concurrency_actions(
-        self, service_name: str
-    ) -> WorkflowConcurrencyOpts | None:
-        if len(self.concurrency_actions) > 0 and self.config.concurrency:
-            raise ValueError(
-                "Error: Both concurrencyActions and concurrency_expression are defined. Please use only one concurrency configuration method."
-            )
-
-        if len(self.concurrency_actions) > 0:
-            action = self.concurrency_actions[0]
-
-            return WorkflowConcurrencyOpts(
-                action=service_name + ":" + action.name,
-                max_runs=action.concurrency__max_runs,
-                limit_strategy=maybe_int_to_str(
-                    convert_python_enum_to_proto(
-                        action.concurrency__limit_strategy,
-                        ConcurrencyLimitStrategyProto,
-                    )
-                ),
-            )
-
-        if self.config.concurrency:
-            return WorkflowConcurrencyOpts(
-                expression=self.config.concurrency.expression,
-                max_runs=self.config.concurrency.max_runs,
-                limit_strategy=self.config.concurrency.limit_strategy,
-            )
-
-        return None
-
-    def validate_on_failure_steps(
-        self, name: str, service_name: str
-    ) -> CreateWorkflowJobOpts | None:
-        if not self.on_failure_steps:
-            return None
-
-        on_failure_step = next(iter(self.on_failure_steps))
-
-        return CreateWorkflowJobOpts(
-            name=name + "-on-failure",
-            steps=[
-                CreateWorkflowStepOpts(
-                    readable_id=on_failure_step.name,
-                    action=service_name + ":" + on_failure_step.name,
-                    timeout=on_failure_step.timeout or "60s",
-                    inputs="{}",
-                    parents=[],
-                    retries=on_failure_step.retries,
-                    rate_limits=on_failure_step.rate_limits,
-                    backoff_factor=on_failure_step.backoff_factor,
-                    backoff_max_seconds=on_failure_step.backoff_max_seconds,
-                )
-            ],
-        )
-
-    def validate_priority(self, default_priority: int | None) -> int | None:
-        validated_priority = (
-            max(1, min(3, default_priority)) if default_priority else None
-        )
-        if validated_priority != default_priority:
-            logger.warning(
-                "Warning: Default Priority Must be between 1 and 3 -- inclusively. Adjusted to be within the range."
-            )
-
-        return validated_priority
-
-    def get_create_opts(self, namespace: str) -> CreateWorkflowVersionOpts:
-        service_name = self.get_service_name(namespace)
-
-        name = self.get_name(namespace)
-        event_triggers = [namespace + event for event in self.config.on_events]
-
-        create_step_opts = [
-            CreateWorkflowStepOpts(
-                readable_id=step.name,
-                action=service_name + ":" + step.name,
-                timeout=step.timeout or "60s",
-                inputs="{}",
-                parents=[x for x in step.parents],
-                retries=step.retries,
-                rate_limits=step.rate_limits,
-                worker_labels=step.desired_worker_labels,
-                backoff_factor=step.backoff_factor,
-                backoff_max_seconds=step.backoff_max_seconds,
-            )
-            for step in self.steps
-            if step.type == StepType.DEFAULT
-        ]
-
-        concurrency = self.validate_concurrency_actions(service_name)
-        on_failure_job = self.validate_on_failure_steps(name, service_name)
-        validated_priority = self.validate_priority(self.config.default_priority)
-
-        return CreateWorkflowVersionOpts(
-            name=name,
-            kind=WorkflowKind.DAG,
-            version=self.config.version,
-            event_triggers=event_triggers,
-            cron_triggers=self.config.on_crons,
-            schedule_timeout=self.config.schedule_timeout,
-            sticky=maybe_int_to_str(
-                convert_python_enum_to_proto(self.config.sticky, StickyStrategyProto)
-            ),
-            jobs=[
-                CreateWorkflowJobOpts(
-                    name=name,
-                    steps=create_step_opts,
-                )
-            ],
-            on_failure_job=on_failure_job,
-            concurrency=concurrency,
-            default_priority=validated_priority,
-        )