dex-python-sdk 0.0.1__py3-none-any.whl

dex/search_attributes.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2022-2026 Super Durable, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Union
+
+from dex.errors import WorkflowDefinitionError
+from dex.dex_api.models import SearchAttribute, SearchAttributeValueType
+from dex.utils.dex_typing import unset_to_none
+
+
+class SearchAttributes:
+    _key_to_type_map: dict[str, SearchAttributeValueType]
+    _int64_attribute_map: dict[str, Union[int, None]]
+    _upsert_to_server_int64_attribute_map: dict[str, Union[int, None]]
+    _string_attribute_map: dict[str, Union[str, None]]
+    _upsert_to_server_string_attribute_map: dict[str, Union[str, None]]
+    _double_attribute_map: dict[str, Union[float, None]]
+    _upsert_to_server_double_attribute_map: dict[str, Union[float, None]]
+    _bool_attribute_map: dict[str, Union[bool, None]]
+    _upsert_to_server_bool_attribute_map: dict[str, Union[bool, None]]
+    _string_array_attribute_map: dict[str, Union[list[str], None]]
+    _upsert_to_server_string_array_attribute_map: dict[str, Union[list[str], None]]
+
+    def __init__(
+        self,
+        key_to_type_map: dict[str, SearchAttributeValueType],
+        search_attribute_map: Union[list[SearchAttribute], None],
+    ):
+        self._key_to_type_map = key_to_type_map
+        self._int64_attribute_map = {}
+        self._upsert_to_server_int64_attribute_map = {}
+        self._string_attribute_map = {}
+        self._upsert_to_server_string_attribute_map = {}
+        self._double_attribute_map = {}
+        self._upsert_to_server_double_attribute_map = {}
+        self._bool_attribute_map = {}
+        self._upsert_to_server_bool_attribute_map = {}
+        self._string_array_attribute_map = {}
+        self._upsert_to_server_string_array_attribute_map = {}
+
+        if search_attribute_map is not None:
+            for attribute in search_attribute_map:
+                key = unset_to_none(attribute.key)
+                if key is None:
+                    raise ValueError("empty search attribute value type is not allowed")
+                value_type = key_to_type_map[key]
+
+                if (
+                    value_type == SearchAttributeValueType.KEYWORD
+                    or value_type == SearchAttributeValueType.DATETIME
+                    or value_type == SearchAttributeValueType.TEXT
+                ):
+                    self._string_attribute_map[key] = unset_to_none(
+                        attribute.string_value
+                    )
+                elif value_type == SearchAttributeValueType.INT:
+                    self._int64_attribute_map[key] = unset_to_none(
+                        attribute.integer_value
+                    )
+                elif value_type == SearchAttributeValueType.DOUBLE:
+                    self._double_attribute_map[key] = unset_to_none(
+                        attribute.double_value
+                    )
+                elif value_type == SearchAttributeValueType.BOOL:
+                    self._bool_attribute_map[key] = unset_to_none(attribute.bool_value)
+                elif value_type == SearchAttributeValueType.KEYWORD_ARRAY:
+                    self._string_array_attribute_map[key] = unset_to_none(
+                        attribute.string_array_value
+                    )
+                else:
+                    raise ValueError(
+                        f"not supported search attribute value type, {value_type}"
+                    )
+
+    def get_search_attribute_int64(self, key: str) -> Union[int, None]:
+        return self._int64_attribute_map.get(key)
+
+    def set_search_attribute_int64(self, key: str, value: Union[int, None]):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.INT
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as int64")
+        self._int64_attribute_map[key] = value
+        self._upsert_to_server_int64_attribute_map[key] = value
+
+    def get_search_attribute_double(self, key: str) -> Union[float, None]:
+        return self._double_attribute_map.get(key)
+
+    def set_search_attribute_double(self, key: str, value: Union[float, None]):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.DOUBLE
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as double")
+        self._double_attribute_map[key] = value
+        self._upsert_to_server_double_attribute_map[key] = value
+
+    def get_search_attribute_boolean(self, key: str) -> Union[bool, None]:
+        return self._bool_attribute_map.get(key)
+
+    def set_search_attribute_boolean(self, key: str, value: Union[bool, None]):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.BOOL
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as bool")
+        self._bool_attribute_map[key] = value
+        self._upsert_to_server_bool_attribute_map[key] = value
+
+    def get_search_attribute_keyword(self, key: str) -> Union[str, None]:
+        return self._string_attribute_map.get(key)
+
+    def set_search_attribute_keyword(self, key: str, value: Union[str, None]):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.KEYWORD
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as keyword")
+        self._string_attribute_map[key] = value
+        self._upsert_to_server_string_attribute_map[key] = value
+
+    def get_search_attribute_text(self, key: str) -> Union[str, None]:
+        return self._string_attribute_map.get(key)
+
+    def set_search_attribute_text(self, key: str, value: Union[str, None]):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.TEXT
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as text")
+        self._string_attribute_map[key] = value
+        self._upsert_to_server_string_attribute_map[key] = value
+
+    def get_search_attribute_datetime(self, key: str) -> Union[str, None]:
+        return self._string_attribute_map.get(key)
+
+    def set_search_attribute_datetime(self, key: str, value: Union[str, None]):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.DATETIME
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as datetime")
+        self._string_attribute_map[key] = value
+        self._upsert_to_server_string_attribute_map[key] = value
+
+    def get_search_attribute_keyword_array(self, key: str) -> Union[list[str], None]:
+        return self._string_array_attribute_map.get(key)
+
+    def set_search_attribute_keyword_array(
+        self, key: str, value: Union[list[str], None]
+    ):
+        if (
+            key not in self._key_to_type_map
+            or self._key_to_type_map[key] != SearchAttributeValueType.KEYWORD_ARRAY
+        ):
+            raise WorkflowDefinitionError(f"key {key} is not defined as keyword array")
+        self._string_array_attribute_map[key] = value
+        self._upsert_to_server_string_array_attribute_map[key] = value
+
+    def get_upsert_to_server_int64_attribute_map(self) -> dict[str, Union[int, None]]:
+        return self._upsert_to_server_int64_attribute_map
+
+    def get_upsert_to_server_string_attribute_map(self) -> dict[str, Union[str, None]]:
+        return self._upsert_to_server_string_attribute_map
+
+    def get_upsert_to_server_string_array_attribute_map(
+        self,
+    ) -> dict[str, Union[list[str], None]]:
+        return self._upsert_to_server_string_array_attribute_map
+
+    def get_upsert_to_server_bool_attribute_map(self) -> dict[str, Union[bool, None]]:
+        return self._upsert_to_server_bool_attribute_map
+
+    def get_upsert_to_server_double_attribute_map(
+        self,
+    ) -> dict[str, Union[float, None]]:
+        return self._upsert_to_server_double_attribute_map

dex/state_decision.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2022-2026 Super Durable, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations  # <-- Additional import.
+
+import typing
+
+from dex.dex_api.models import WorkflowConditionalClose, WorkflowConditionalCloseType
+from dex.workflow_state_options import WorkflowStateOptions
+
+if typing.TYPE_CHECKING:
+    from dex.registry import Registry
+    from dex.workflow_state import WorkflowState
+
+from dataclasses import dataclass
+from typing import Any, List, Union, Optional
+
+from dex.dex_api.models.state_decision import StateDecision as IdlStateDecision
+
+from dex.object_encoder import ObjectEncoder
+
+from dex.state_movement import StateMovement, _to_idl_state_movement
+
+
+@dataclass
+class InternalConditionalClose:
+    conditional_close_type: WorkflowConditionalCloseType
+
+    channel_name: str
+
+    close_input: Any = None
+
+
+@dataclass
+class StateDecision:
+    next_states: List[StateMovement]
+
+    conditional_close: typing.Optional[InternalConditionalClose] = None
+
+    dead_end: typing.ClassVar[StateDecision]
+
+    @classmethod
+    def graceful_complete_workflow(cls, output: Any = None) -> StateDecision:
+        return StateDecision([StateMovement.graceful_complete_workflow(output)])
+
+    @classmethod
+    def force_complete_workflow(cls, output: Any = None) -> StateDecision:
+        return StateDecision([StateMovement.force_complete_workflow(output)])
+
+    @classmethod
+    def force_fail_workflow(cls, output: Any = None) -> StateDecision:
+        return StateDecision([StateMovement.force_fail_workflow(output)])
+
+    @classmethod
+    def single_next_state(
+        cls,
+        state: Union[str, type[WorkflowState]],
+        state_input: Any = None,
+        state_options_override: Optional[WorkflowStateOptions] = None,
+        wait_for_key: Optional[str] = None,
+    ) -> StateDecision:
+        return StateDecision(
+            [
+                StateMovement.create(
+                    state, state_input, state_options_override, wait_for_key
+                )
+            ]
+        )
+
+    @classmethod
+    def multi_next_states(
+        cls, *next_states: Union[type[WorkflowState], StateMovement]
+    ) -> StateDecision:
+        next_list = [
+            n if isinstance(n, StateMovement) else StateMovement.create(n)
+            for n in next_states
+        ]
+        return StateDecision(next_list)
+
+    # Atomically force complete the workflow if internal channel is empty, otherwise trigger the state movements from the current thread
+    # This is to ensure all the messages in the channel are processed before completing the workflow, otherwise messages may be lost.
+    # Without this atomic API, if just checking the channel emptiness in the State WaitUntil, a workflow may receive new messages during the
+    # execution of state APIs.
+    #
+    # Note that it's only for internal messages published from RPCs.
+    # It doesn't cover the cases that internal messages are published from other State APIs.
+    # If you do want to use other State APIs to publish messages to the channel at the same time, you can use persistence locking to
+    # ensure only the State APIs are not executed in parallel.
+    @classmethod
+    def force_complete_if_internal_channel_empty_or_else(
+        cls,
+        internal_channel_name: str,
+        workflow_complete_output: Any = None,  # if channel is empty, complete the workflow with the output
+        or_else_state: Union[
+            str, type[WorkflowState]
+        ] = "",  # required not empty -- if channel is NOT empty, go to this state with the state input
+        state_input: Any = None,
+    ) -> StateDecision:
+        return StateDecision(
+            [StateMovement.create(or_else_state, state_input)],
+            InternalConditionalClose(
+                WorkflowConditionalCloseType.FORCE_COMPLETE_ON_INTERNAL_CHANNEL_EMPTY,
+                internal_channel_name,
+                workflow_complete_output,
+            ),
+        )
+
+    # Atomically force complete the workflow if signal channel is empty, otherwise trigger the state movements from the current thread
+    # This is to ensure all the messages in the channel are processed before completing the workflow, otherwise messages may be lost.
+    # Without this atomic API, if just checking the channel emptiness in the State WaitUntil, a workflow may receive new messages during the
+    # execution of state APIs.
+    @classmethod
+    def force_complete_if_signal_channel_empty_or_else(
+        cls,
+        signal_channel_name: str,
+        workflow_complete_output: Any = None,  # if channel is empty, complete the workflow with the output
+        or_else_state: Union[
+            str, type[WorkflowState]
+        ] = "",  # required not empty-- if channel is NOT empty, go to this state with the state input
+        state_input: Any = None,
+    ) -> StateDecision:
+        return StateDecision(
+            [StateMovement.create(or_else_state, state_input)],
+            InternalConditionalClose(
+                WorkflowConditionalCloseType.FORCE_COMPLETE_ON_SIGNAL_CHANNEL_EMPTY,
+                signal_channel_name,
+                workflow_complete_output,
+            ),
+        )
+
+
+StateDecision.dead_end = StateDecision([StateMovement.dead_end])
+
+
+def _to_idl_state_decision(
+    decision: StateDecision, wf_type: str, registry: Registry, encoder: ObjectEncoder
+) -> IdlStateDecision:
+    idl_decision = IdlStateDecision()
+    if len(decision.next_states) > 0:
+        idl_decision.next_states = [
+            _to_idl_state_movement(movement, wf_type, registry, encoder)
+            for movement in decision.next_states
+        ]
+    if decision.conditional_close is not None:
+        idl_decision.conditional_close = WorkflowConditionalClose(
+            conditional_close_type=decision.conditional_close.conditional_close_type,
+            channel_name=decision.conditional_close.channel_name,
+            close_input=encoder.encode(decision.conditional_close.close_input),
+        )
+    return idl_decision

dex/state_execution_locals.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2022-2026 Super Durable, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List, Tuple, Union
+
+from dex.errors import WorkflowDefinitionError
+from dex.dex_api.models import EncodedObject, KeyValue
+from dex.dex_api.types import Unset
+from dex.object_encoder import ObjectEncoder
+
+
+class StateExecutionLocals:
+    _record_events: dict[str, Union[EncodedObject, Unset]]
+    _attribute_name_to_encoded_object_map: dict[str, Union[EncodedObject, Unset]]
+    _upsert_attributes_to_return_to_server: dict[str, Union[EncodedObject, Unset]]
+    _object_encoder: ObjectEncoder
+
+    def __init__(
+        self,
+        attribute_name_to_encoded_object_map: dict[str, Union[EncodedObject, Unset]],
+        object_encoder: ObjectEncoder,
+    ):
+        self._object_encoder = object_encoder
+        self._attribute_name_to_encoded_object_map = (
+            attribute_name_to_encoded_object_map
+        )
+        self._upsert_attributes_to_return_to_server = {}
+        self._record_events = {}
+
+    def set_state_execution_local(self, key: str, value: Any):
+        encoded_data = self._object_encoder.encode(value)
+        self._attribute_name_to_encoded_object_map[key] = encoded_data
+        self._upsert_attributes_to_return_to_server[key] = encoded_data
+
+    def get_state_execution_local(self, key: str) -> Any:
+        encoded_object = self._attribute_name_to_encoded_object_map.get(key)
+        if encoded_object is None:
+            return None
+        return self._object_encoder.decode(encoded_object)
+
+    def record_event(self, key: str, *event_data: Tuple[Any, ...]):
+        if key in self._record_events:
+            raise WorkflowDefinitionError("Cannot record the same event more than once")
+
+        if event_data is not None and len(event_data) == 1:
+            self._record_events[key] = self._object_encoder.encode(event_data[0])
+
+        self._record_events[key] = self._object_encoder.encode(event_data)
+
+    def get_upsert_state_execution_local_attributes(self) -> List[KeyValue]:
+        return [
+            KeyValue(item_key, item_value)
+            for item_key, item_value in self._upsert_attributes_to_return_to_server.items()
+        ]
+
+    def get_record_events(self) -> List[KeyValue]:
+        return [
+            KeyValue(item_key, item_value)
+            for item_key, item_value in self._record_events.items()
+        ]

dex/state_movement.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2022-2026 Super Durable, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import typing
+from typing import Union, Optional
+
+from dex.errors import WorkflowDefinitionError
+
+if typing.TYPE_CHECKING:
+    from dex.workflow_state import (
+        WorkflowState,
+    )
+    from dex.registry import Registry
+
+from dataclasses import dataclass
+from typing import Any
+
+
+from dex.dex_api.models.state_movement import StateMovement as IdlStateMovement
+
+from dex.object_encoder import ObjectEncoder
+
+from dex.workflow_state_options import _to_idl_state_options, WorkflowStateOptions
+
+reserved_state_id_prefix = "_SYS_"
+
+graceful_completing_sys_state_id = (
+    reserved_state_id_prefix + "GRACEFUL_COMPLETING_WORKFLOW"
+)
+force_completing_sys_state_id = reserved_state_id_prefix + "FORCE_COMPLETING_WORKFLOW"
+force_failing_sys_state_id = reserved_state_id_prefix + "FORCE_FAILING_WORKFLOW"
+dead_end_sys_state_id = reserved_state_id_prefix + "DEAD_END"
+
+
+@dataclass
+class StateMovement:
+    state_id: str
+    state_input: Any = None
+    state_options_override: Optional[WorkflowStateOptions] = None
+    wait_for_key: Optional[str] = None
+
+    dead_end: typing.ClassVar[StateMovement]
+
+    @classmethod
+    def graceful_complete_workflow(cls, output: Any = None) -> StateMovement:
+        return StateMovement(graceful_completing_sys_state_id, output)
+
+    @classmethod
+    def force_complete_workflow(cls, output: Any = None) -> StateMovement:
+        return StateMovement(force_completing_sys_state_id, output)
+
+    @classmethod
+    def force_fail_workflow(cls, output: Any = None) -> StateMovement:
+        return StateMovement(force_failing_sys_state_id, output)
+
+    @classmethod
+    def create(
+        cls,
+        state: Union[str, type[WorkflowState]],
+        state_input: Any = None,
+        state_options_override: Optional[WorkflowStateOptions] = None,
+        wait_for_key: Optional[str] = None,
+    ) -> StateMovement:
+        if isinstance(state, str):
+            state_id = state
+        else:
+            from dex.workflow_state import (
+                get_state_id_by_class,
+            )
+
+            state_id = get_state_id_by_class(state)
+        if state_id.startswith(reserved_state_id_prefix):
+            raise WorkflowDefinitionError("cannot use reserved stateId")
+        return StateMovement(
+            state_id, state_input, state_options_override, wait_for_key
+        )
+
+
+StateMovement.dead_end = StateMovement(dead_end_sys_state_id)
+
+
+def _to_idl_state_movement(
+    movement: StateMovement, wf_type: str, registry: Registry, encoder: ObjectEncoder
+) -> IdlStateMovement:
+    idl_movement = IdlStateMovement(
+        state_id=movement.state_id, state_input=encoder.encode(movement.state_input)
+    )
+    if not movement.state_id.startswith(reserved_state_id_prefix):
+        state = registry.get_workflow_state_with_check(wf_type, movement.state_id)
+
+        from dex.workflow_state import (
+            should_skip_wait_until,
+        )
+
+        if movement.state_options_override is not None:
+            options = movement.state_options_override
+        else:
+            options = state.get_state_options()
+
+        idl_state_options = _to_idl_state_options(
+            should_skip_wait_until(state),
+            options,
+            registry.get_state_store(wf_type),
+        )
+
+        idl_movement.state_options = idl_state_options
+
+        if movement.wait_for_key is not None:
+            idl_movement.wait_for_key = movement.wait_for_key
+    return idl_movement

dex/state_schema.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2022-2026 Super Durable, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import List
+
+from dex.workflow_state import WorkflowState
+
+
+@dataclass
+class StateDef:
+    state: WorkflowState
+    can_start_workflow: bool
+
+    @classmethod
+    def starting_state(cls, state: WorkflowState):
+        return StateDef(state, True)
+
+    @classmethod
+    def non_starting_state(cls, state: WorkflowState):
+        return StateDef(state, False)
+
+
+@dataclass
+class StateSchema:
+    states: List[StateDef] = field(default_factory=list)
+
+    # TODO: it's super weird that we can't use type hint here " ->StateSchema" for return
+    # But the pattern works for state_movement.py
+    @classmethod
+    def with_starting_state(
+        cls, starting_state: WorkflowState, *non_starting_states: WorkflowState
+    ):
+        return StateSchema(
+            [StateDef.starting_state(starting_state)]
+            + [StateDef.non_starting_state(s) for s in non_starting_states]
+        )
+
+    @classmethod
+    def no_starting_state(cls, *non_starting_states: WorkflowState):
+        return StateSchema(
+            [StateDef.non_starting_state(s) for s in non_starting_states]
+        )

dex/stop_workflow_options.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2022-2026 Super Durable, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from dex.dex_api.models import WorkflowStopType
+
+
+@dataclass
+class StopWorkflowOptions:
+    workflow_stop_type: WorkflowStopType
+    reason: str