IWF-681 add data attributes API methods (#94)

* IWF-681 add missing data attributes API methods

---------

Co-authored-by: Nathan Givens <[email protected]>

Author: N-give

iwf/client.py

@@ -1,10 +1,10 @@
 import inspect
-from typing import Any, Callable, Optional, Type, TypeVar, Union
+from typing import Any, Callable, List, Optional, Type, TypeVar, Union
 
 from typing_extensions import deprecated
 
 from iwf.client_options import ClientOptions
-from iwf.errors import InvalidArgumentError, WorkflowDefinitionError
+from iwf.errors import InvalidArgumentError, NotRegisteredError, WorkflowDefinitionError
 from iwf.iwf_api.models import (
     SearchAttribute,
     SearchAttributeKeyAndType,
@@ -162,6 +162,71 @@ def stop_workflow(
     ):
         return self._unregistered_client.stop_workflow(workflow_id, "", options)
 
+    def get_all_workflow_data_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        workflow_run_id: str = "",
+    ):
+        return self.get_workflow_data_attributes(
+            workflow_class, workflow_id, workflow_run_id, None
+        )
+
+    def get_workflow_data_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        workflow_run_id: str = "",
+        keys: Optional[List[str]] = None,
+    ):
+        wf_type = get_workflow_type_by_class(workflow_class)
+        data_attr_type_store = self._registry.get_data_attribute_types(wf_type)
+        if keys:
+            for key in keys:
+                if not data_attr_type_store.is_valid_name_or_prefix(key):
+                    raise NotRegisteredError(
+                        f"key {key} is not registered in workflow {wf_type}"
+                    )
+
+        response = self._unregistered_client.get_workflow_data_attributes(
+            workflow_id, workflow_run_id, keys
+        )
+
+        if not response.objects:
+            raise RuntimeError("data attributes not returned")
+
+        res = {}
+        for kv in response.objects:
+            k = unset_to_none(kv.key)
+            if k and kv.value:
+                res[kv.key] = self._options.object_encoder.decode(
+                    kv.value, data_attr_type_store.get_type(k)
+                )
+
+        return res
+
+    def set_workflow_data_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        workflow_run_id: str = "",
+        data_attributes: dict[str, Any] = dict(),
+    ):
+        wf_type = get_workflow_type_by_class(workflow_class)
+        data_attr_type_store = self._registry.get_data_attribute_types(wf_type)
+        for key, value in data_attributes.items():
+            if not data_attr_type_store.is_valid_name_or_prefix(key):
+                raise NotRegisteredError(f"data attribute {key} is not registered")
+
+            data_attr_type = data_attr_type_store.get_type(key)
+            if not isinstance(value, data_attr_type):
+                raise NotRegisteredError(
+                    f"data attribute {key} is not registered as {type(value)}"
+                )
+        return self._unregistered_client.set_workflow_data_attributes(
+            workflow_id, workflow_run_id, data_attributes
+        )
+
     def invoke_rpc(
         self,
         workflow_id: str,

iwf/data_attributes.py

@@ -1,46 +1,46 @@
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 from iwf.errors import WorkflowDefinitionError
 from iwf.iwf_api.models import EncodedObject
 from iwf.object_encoder import ObjectEncoder
+from iwf.type_store import TypeStore
 
 
 class DataAttributes:
-    _type_store: dict[str, Optional[type]]
-    _prefix_type_store: dict[str, Optional[type]]
+    _type_store: TypeStore
     _object_encoder: ObjectEncoder
     _current_values: dict[str, Union[EncodedObject, None]]
     _updated_values_to_return: dict[str, EncodedObject]
 
     def __init__(
         self,
-        type_store: dict[str, Optional[type]],
-        prefix_type_store: dict[str, Optional[type]],
+        type_store: TypeStore,
         object_encoder: ObjectEncoder,
         current_values: dict[str, Union[EncodedObject, None]],
     ):
         self._object_encoder = object_encoder
         self._type_store = type_store
-        self._prefix_type_store = prefix_type_store
         self._current_values = current_values
         self._updated_values_to_return = {}
 
     def get_data_attribute(self, key: str) -> Any:
-        is_registered, registered_type = self._validate_key_and_get_type(key)
+        is_registered = self._type_store.is_valid_name_or_prefix(key)
         if not is_registered:
             raise WorkflowDefinitionError(f"data attribute %s is not registered {key}")
 
         encoded_object = self._current_values.get(key)
         if encoded_object is None:
             return None
 
+        registered_type = self._type_store.get_type(key)
         return self._object_encoder.decode(encoded_object, registered_type)
 
     def set_data_attribute(self, key: str, value: Any):
-        is_registered, registered_type = self._validate_key_and_get_type(key)
+        is_registered = self._type_store.is_valid_name_or_prefix(key)
         if not is_registered:
             raise WorkflowDefinitionError(f"data attribute %s is not registered {key}")
 
+        registered_type = self._type_store.get_type(key)
         if registered_type is not None and not isinstance(value, registered_type):
             raise WorkflowDefinitionError(
                 f"data attribute %s is of the right type {registered_type}"
@@ -52,13 +52,3 @@ def set_data_attribute(self, key: str, value: Any):
 
     def get_updated_values_to_return(self) -> dict[str, EncodedObject]:
         return self._updated_values_to_return
-
-    def _validate_key_and_get_type(self, key) -> tuple[bool, Optional[type]]:
-        if key in self._type_store:
-            return (True, self._type_store.get(key))
-
-        for prefix, t in self._prefix_type_store.items():
-            if key.startswith(prefix):
-                return (True, t)
-
-        return (False, None)

iwf/registry.py

@@ -16,8 +16,7 @@ class Registry:
     _state_store: dict[str, dict[str, WorkflowState]]
     _internal_channel_type_store: dict[str, TypeStore]
     _signal_channel_type_store: dict[str, dict[str, Optional[type]]]
-    _data_attribute_types: dict[str, dict[str, Optional[type]]]
-    _data_attribute_prefix_types: dict[str, dict[str, Optional[type]]]
+    _data_attribute_types: dict[str, TypeStore]
     _search_attribute_types: dict[str, dict[str, SearchAttributeValueType]]
     _rpc_infos: dict[str, dict[str, RPCInfo]]
 
@@ -28,7 +27,6 @@ def __init__(self):
         self._internal_channel_type_store = dict()
         self._signal_channel_type_store = dict()
         self._data_attribute_types = dict()
-        self._data_attribute_prefix_types = dict()
         self._search_attribute_types = {}
         self._rpc_infos = dict()
 
@@ -38,7 +36,6 @@ def add_workflow(self, wf: ObjectWorkflow):
         self._register_internal_channels(wf)
         self._register_signal_channels(wf)
         self._register_data_attributes(wf)
-        self._register_data_attribute_prefix_types(wf)
         self._register_search_attributes(wf)
         self._register_workflow_rpcs(wf)
 
@@ -77,19 +74,14 @@ def get_internal_channel_type_store(self, wf_type: str) -> TypeStore:
     def get_signal_channel_types(self, wf_type: str) -> dict[str, Optional[type]]:
         return self._signal_channel_type_store[wf_type]
 
-    def get_data_attribute_types(self, wf_type: str) -> dict[str, Optional[type]]:
+    def get_data_attribute_types(self, wf_type: str) -> TypeStore:
         return self._data_attribute_types[wf_type]
 
     def get_search_attribute_types(
         self, wf_type: str
     ) -> dict[str, SearchAttributeValueType]:
         return self._search_attribute_types[wf_type]
 
-    def get_data_attribute_prefix_types(
-        self, wf_type: str
-    ) -> dict[str, Optional[type]]:
-        return self._data_attribute_prefix_types[wf_type]
-
     def get_rpc_infos(self, wf_type: str) -> dict[str, RPCInfo]:
         return self._rpc_infos[wf_type]
 
@@ -123,10 +115,13 @@ def _register_signal_channels(self, wf: ObjectWorkflow):
 
     def _register_data_attributes(self, wf: ObjectWorkflow):
         wf_type = get_workflow_type(wf)
-        data_attribute_types: dict[str, Optional[type]] = {}
+        data_attribute_types: TypeStore = TypeStore(Type.DATA_ATTRIBUTE)
         for field in wf.get_persistence_schema().persistence_fields:
-            if field.field_type == PersistenceFieldType.DataAttribute:
-                data_attribute_types[field.key] = field.value_type
+            if (
+                field.field_type == PersistenceFieldType.DataAttribute
+                or field.field_type == PersistenceFieldType.DataAttributePrefix
+            ):
+                data_attribute_types.add_data_attribute_def(field)
         self._data_attribute_types[wf_type] = data_attribute_types
 
     def _register_search_attributes(self, wf: ObjectWorkflow):
@@ -146,14 +141,6 @@ def _register_search_attributes(self, wf: ObjectWorkflow):
                 types[field.key] = sa_type
         self._search_attribute_types[wf_type] = types
 
-    def _register_data_attribute_prefix_types(self, wf: ObjectWorkflow):
-        wf_type = get_workflow_type(wf)
-        data_attribute_prefix_types: dict[str, Optional[type]] = {}
-        for field in wf.get_persistence_schema().persistence_fields:
-            if field.field_type == PersistenceFieldType.DataAttributePrefix:
-                data_attribute_prefix_types[field.key] = field.value_type
-        self._data_attribute_prefix_types[wf_type] = data_attribute_prefix_types
-
     def _register_workflow_state(self, wf):
         wf_type = get_workflow_type(wf)
         state_map = {}

iwf/tests/test_persistence_data_attributes.py

@@ -15,7 +15,6 @@
 from iwf.workflow_context import WorkflowContext
 from iwf.workflow_options import WorkflowOptions
 from iwf.workflow_state import T, WorkflowState
-from iwf.rpc import rpc
 
 initial_da_1 = "initial_da_1"
 initial_da_value_1 = "value_1"
@@ -35,6 +34,18 @@
 test_da_prefix_key_2 = "test-da-prefix-2"
 test_da_prefix_value_1 = "test-da-value-1"
 test_da_prefix_value_2 = "test-da-value-2"
+test_da_set_key = "test_da_set_key"
+test_da_set_value = "test_da_set_value"
+
+expected_final_das = {
+    initial_da_1: initial_da_value_1,
+    initial_da_2: final_initial_da_value_2,
+    test_da_1: final_test_da_value_1,
+    test_da_2: final_test_da_value_2,
+    test_da_prefix_key_1: test_da_prefix_value_1,
+    test_da_prefix_key_2: test_da_prefix_value_2,
+    test_da_set_key: test_da_set_value,
+}
 
 
 class DataAttributeRWState(WorkflowState[None]):
@@ -82,17 +93,7 @@ def get_persistence_schema(self) -> PersistenceSchema:
             PersistenceField.data_attribute_def(test_da_1, str),
             PersistenceField.data_attribute_def(test_da_2, int),
             PersistenceField.data_attribute_prefix_def(test_da_prefix, str),
-        )
-
-    @rpc()
-    def test_persistence_read(self, pers: Persistence):
-        return (
-            pers.get_data_attribute(initial_da_1),
-            pers.get_data_attribute(initial_da_2),
-            pers.get_data_attribute(test_da_1),
-            pers.get_data_attribute(test_da_2),
-            pers.get_data_attribute(test_da_prefix_key_1),
-            pers.get_data_attribute(test_da_prefix_key_2),
+            PersistenceField.data_attribute_def(test_da_set_key, str),
         )
 
 
@@ -113,19 +114,22 @@ def test_persistence_data_attributes_workflow(self):
             },
         )
 
-        self.client.start_workflow(
+        wf_run_id = self.client.start_workflow(
             PersistenceDataAttributesWorkflow, wf_id, 100, None, start_options
         )
+        self.client.set_workflow_data_attributes(
+            PersistenceDataAttributesWorkflow,
+            wf_id,
+            wf_run_id,
+            {
+                test_da_set_key: test_da_set_value,
+            },
+        )
         self.client.wait_for_workflow_completion(wf_id, None)
 
-        res = self.client.invoke_rpc(
-            wf_id, PersistenceDataAttributesWorkflow.test_persistence_read
+        all_data_attributes = self.client.get_all_workflow_data_attributes(
+            PersistenceDataAttributesWorkflow, wf_id, wf_run_id
         )
-        assert res == [
-            final_initial_da_value_1,
-            final_initial_da_value_2,
-            final_test_da_value_1,
-            final_test_da_value_2,
-            test_da_prefix_value_1,
-            test_da_prefix_value_2,
-        ]
+        for k, v in expected_final_das.items():
+            assert k in all_data_attributes
+            assert all_data_attributes[k] == v

iwf/type_store.py

@@ -3,12 +3,13 @@
 
 from iwf.communication_schema import CommunicationMethod
 from iwf.errors import WorkflowDefinitionError, NotRegisteredError
+from iwf.persistence_schema import PersistenceField, PersistenceFieldType
 
 
 class Type(Enum):
     INTERNAL_CHANNEL = 1
+    DATA_ATTRIBUTE = 2
     # TODO: extend to other types
-    # DATA_ATTRIBUTE = 2
     # SIGNAL_CHANNEL = 3
 
 
@@ -44,6 +45,17 @@ def add_internal_channel_def(self, obj: CommunicationMethod):
             )
         self._do_add_to_store(obj.is_prefix, obj.name, obj.value_type)
 
+    def add_data_attribute_def(self, obj: PersistenceField):
+        if self._class_type != Type.DATA_ATTRIBUTE:
+            raise ValueError(
+                f"Cannot add internal channel definition to {self._class_type}"
+            )
+        self._do_add_to_store(
+            obj.field_type == PersistenceFieldType.DataAttributePrefix,
+            obj.key,
+            obj.value_type,
+        )
+
     def _validate_name(self, name: str) -> bool:
         if name in self._name_to_type_store:
             return True