IWF-628: Add setting search attributes (#84)

lwolczynski · web-flow · commit 622edae534c8 · 2025-03-17T12:27:02.000-05:00
* IWF-628: Add set search attributes methods

* IWF-628: Add unit tests

* IWF-628: Fix

* IWF-628: Fix test

* IWF-628: Fix

* IWF-628: Testing

* IWF-628: Fix test

* IWF-628: Fix test

* IWF-628: Fix test

* IWF-628: Add comment

* IWF-628: Change test

* IWF-628: Lint

* IWF-628: Simplify test

* IWF-628: Refactor

* IWF-628: Change comment

* IWF-628: Change timeout

* IWF-628: Testing

* IWF-628: Unskip tests

* IWF-628: Edit comment

* IWF-628: Increase wait
diff --git a/.devcontainer/.env b/.devcontainer/.env
@@ -2,5 +2,6 @@ CASSANDRA_VERSION=3.11.9
 ELASTICSEARCH_VERSION=7.16.2
 MYSQL_VERSION=8
 POSTGRESQL_VERSION=13
-TEMPORAL_VERSION=1.20.2
-TEMPORAL_UI_VERSION=2.15.0
+TEMPORAL_VERSION=1.25
+TEMPORAL_ADMIN_TOOLS_VERSION=1.25.2-tctl-1.18.1-cli-1.1.1
+TEMPORAL_UI_VERSION=2.31.2
diff --git a/iwf/client.py b/iwf/client.py
@@ -225,8 +225,10 @@ def get_all_search_attributes(
         workflow_id: str,
         workflow_run_id: Optional[str] = None,
     ):
+        run_id = workflow_run_id if workflow_run_id is not None else ""
+
         return self._do_get_workflow_search_attributes(
-            workflow_class, workflow_id, workflow_run_id
+            workflow_class, workflow_id, run_id
         )
 
     def get_workflow_search_attributes(
@@ -240,24 +242,27 @@ def get_workflow_search_attributes(
             raise ValueError(
                 "attribute_keys must contain at least one entry, or use get_all_search_attributes API to get all"
             )
+
+        run_id = workflow_run_id if workflow_run_id is not None else ""
+
         return self._do_get_workflow_search_attributes(
-            workflow_class, workflow_id, workflow_run_id, attribute_keys
+            workflow_class, workflow_id, run_id, attribute_keys
         )
 
     def _do_get_workflow_search_attributes(
         self,
         workflow_class: type[ObjectWorkflow],
         workflow_id: str,
-        workflow_run_id: Optional[str],
+        workflow_run_id: str,
         attribute_keys: Optional[list[str]] = None,
     ):
         wf_type = get_workflow_type_by_class(workflow_class)
         self._registry.get_workflow_with_check(wf_type)
 
         search_attribute_types = self._registry.get_search_attribute_types(wf_type)
 
-        # if attribute keys is None or empty, iwf server will return all search attributes
-        if attribute_keys is not None and attribute_keys:
+        # if attribute keys is None, will fetch all registered search attributes from the server
+        if attribute_keys:
             non_existing_search_attribute_list: list[str] = []
             for attribute_key in attribute_keys:
                 if attribute_key not in search_attribute_types:
@@ -277,10 +282,8 @@ def _do_get_workflow_search_attributes(
                 sa_type = search_attribute_types[attribute_key]
                 key_and_types.append(SearchAttributeKeyAndType(attribute_key, sa_type))
 
-        run_id = workflow_run_id if workflow_run_id is not None else ""
-
         response = self._unregistered_client.get_workflow_search_attributes(
-            workflow_id, run_id, key_and_types
+            workflow_id, workflow_run_id, key_and_types
         )
 
         response_sas = response.search_attributes
@@ -296,28 +299,78 @@ def _do_get_workflow_search_attributes(
             if response_sa_key is None:
                 raise RuntimeError("search attribute key is None")
             response_sa_type = search_attribute_types[response_sa_key]
-            value = self.get_search_attribute_value(response_sa_type, response_sa)
+            value = get_search_attribute_value(response_sa_type, response_sa)
             result[response_sa_key] = value
 
         return result
 
-    @staticmethod
-    def get_search_attribute_value(
-        sa_type: SearchAttributeValueType, attribute: SearchAttribute
+    def set_workflow_search_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        search_attributes: list[SearchAttribute],
+        workflow_run_id: Optional[str] = None,
     ):
-        if (
-            sa_type == SearchAttributeValueType.KEYWORD
-            or sa_type == SearchAttributeValueType.DATETIME
-            or sa_type == SearchAttributeValueType.TEXT
-        ):
-            return unset_to_none(attribute.string_value)
-        elif sa_type == SearchAttributeValueType.INT:
-            return unset_to_none(attribute.integer_value)
-        elif sa_type == SearchAttributeValueType.DOUBLE:
-            return unset_to_none(attribute.double_value)
-        elif sa_type == SearchAttributeValueType.BOOL:
-            return unset_to_none(attribute.bool_value)
-        elif sa_type == SearchAttributeValueType.KEYWORD_ARRAY:
-            return unset_to_none(attribute.string_array_value)
-        else:
-            raise ValueError(f"not supported search attribute value type, {sa_type}")
+        run_id = workflow_run_id if workflow_run_id is not None else ""
+
+        return self._do_set_workflow_search_attributes(
+            workflow_class, workflow_id, run_id, search_attributes
+        )
+
+    def _do_set_workflow_search_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        workflow_run_id: str,
+        search_attributes: list[SearchAttribute],
+    ):
+        wf_type = get_workflow_type_by_class(workflow_class)
+        self._registry.get_workflow_with_check(wf_type)
+
+        search_attribute_types = self._registry.get_search_attribute_types(wf_type)
+
+        # Check that the requested sa type is registered to the key
+        for search_attribute in search_attributes:
+            sa_key = unset_to_none(search_attribute.key)
+            if sa_key is None:
+                raise RuntimeError("search attribute key is None")
+            if sa_key not in search_attribute_types:
+                raise InvalidArgumentError(f"Search attribute not registered: {sa_key}")
+            registered_value_type = search_attribute_types[sa_key]
+
+            sa_value_type = unset_to_none(search_attribute.value_type)
+            if sa_value_type is None:
+                raise RuntimeError("search value type is None")
+
+            if (
+                sa_value_type is not None
+                and registered_value_type != sa_value_type.value
+            ):
+                raise ValueError(
+                    f"Search attribute key, {sa_key} is registered to type {registered_value_type}, but tried to add search attribute type {sa_value_type.value}"
+                )
+
+        self._unregistered_client.set_workflow_search_attributes(
+            workflow_id, workflow_run_id, search_attributes
+        )
+
+
+def get_search_attribute_value(
+    sa_type: SearchAttributeValueType, attribute: SearchAttribute
+):
+    if (
+        sa_type == SearchAttributeValueType.KEYWORD
+        or sa_type == SearchAttributeValueType.DATETIME
+        or sa_type == SearchAttributeValueType.TEXT
+    ):
+        return unset_to_none(attribute.string_value)
+    elif sa_type == SearchAttributeValueType.INT:
+        return unset_to_none(attribute.integer_value)
+    elif sa_type == SearchAttributeValueType.DOUBLE:
+        return unset_to_none(attribute.double_value)
+    elif sa_type == SearchAttributeValueType.BOOL:
+        return unset_to_none(attribute.bool_value)
+    elif sa_type == SearchAttributeValueType.KEYWORD_ARRAY:
+        return unset_to_none(attribute.string_array_value)
+    else:
+        raise ValueError(f"not supported search attribute value type, {sa_type}")
diff --git a/iwf/persistence.py b/iwf/persistence.py
@@ -25,41 +25,43 @@ def set_data_attribute(self, key: str, value: Any):
     def get_search_attribute_int64(self, key: str) -> Union[None, int]:
         return self._search_attributes.get_search_attribute_int64(key)
 
-    def set_search_attribute_int64(self, key: str, value: int):
+    def set_search_attribute_int64(self, key: str, value: Union[None, int]):
         self._search_attributes.set_search_attribute_int64(key, value)
 
     def get_search_attribute_double(self, key: str) -> Union[None, float]:
         return self._search_attributes.get_search_attribute_double(key)
 
-    def set_search_attribute_double(self, key: str, value: float):
+    def set_search_attribute_double(self, key: str, value: Union[None, float]):
         self._search_attributes.set_search_attribute_double(key, value)
 
     def get_search_attribute_boolean(self, key: str) -> Union[None, bool]:
         return self._search_attributes.get_search_attribute_boolean(key)
 
-    def set_search_attribute_boolean(self, key: str, value: bool):
+    def set_search_attribute_boolean(self, key: str, value: Union[None, bool]):
         self._search_attributes.set_search_attribute_boolean(key, value)
 
     def get_search_attribute_keyword(self, key: str) -> Union[None, str]:
         return self._search_attributes.get_search_attribute_keyword(key)
 
-    def set_search_attribute_keyword(self, key: str, value: str):
+    def set_search_attribute_keyword(self, key: str, value: Union[None, str]):
         self._search_attributes.set_search_attribute_keyword(key, value)
 
     def get_search_attribute_text(self, key: str) -> Union[None, str]:
         return self._search_attributes.get_search_attribute_text(key)
 
-    def set_search_attribute_text(self, key: str, value: str):
+    def set_search_attribute_text(self, key: str, value: Union[None, str]):
         self._search_attributes.set_search_attribute_text(key, value)
 
     def get_search_attribute_datetime(self, key: str) -> Union[None, str]:
         return self._search_attributes.get_search_attribute_datetime(key)
 
-    def set_search_attribute_datetime(self, key: str, value: str):
+    def set_search_attribute_datetime(self, key: str, value: Union[None, str]):
         self._search_attributes.set_search_attribute_datetime(key, value)
 
     def get_search_attribute_keyword_array(self, key: str) -> Union[None, list[str]]:
         return self._search_attributes.get_search_attribute_keyword_array(key)
 
-    def set_search_attribute_keyword_array(self, key: str, value: list[str]):
+    def set_search_attribute_keyword_array(
+        self, key: str, value: Union[None, list[str]]
+    ):
         self._search_attributes.set_search_attribute_keyword_array(key, value)
diff --git a/iwf/search_attributes.py b/iwf/search_attributes.py
@@ -72,7 +72,7 @@ def __init__(
     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: int):
+    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
@@ -84,7 +84,7 @@ def set_search_attribute_int64(self, key: str, value: int):
     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: float):
+    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
@@ -96,7 +96,7 @@ def set_search_attribute_double(self, key: str, value: float):
     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: bool):
+    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
@@ -108,7 +108,7 @@ def set_search_attribute_boolean(self, key: str, value: bool):
     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: str):
+    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
@@ -120,7 +120,7 @@ def set_search_attribute_keyword(self, key: str, value: str):
     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: str):
+    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
@@ -132,7 +132,7 @@ def set_search_attribute_text(self, key: str, value: str):
     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: str):
+    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
@@ -144,7 +144,9 @@ def set_search_attribute_datetime(self, key: str, value: str):
     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: list[str]):
+    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
diff --git a/iwf/tests/iwf-service-env/docker-compose-init.sh b/iwf/tests/iwf-service-env/docker-compose-init.sh
@@ -30,8 +30,10 @@ for run in {1..120}; do
   sleep 0.1
   temporal  operator search-attribute  create --name CustomStringField --type Text
   sleep 0.1
+  temporal  operator search-attribute  create --name CustomKeywordArrayField --type KeywordList
+  sleep 0.1
 
-  if checkExists "IwfWorkflowType" && checkExists "IwfGlobalWorkflowVersion" && checkExists "IwfExecutingStateIds" && checkExists "CustomKeywordField" && checkExists "CustomIntField" && checkExists "CustomBoolField" && checkExists "CustomDoubleField" && checkExists "CustomDatetimeField" && checkExists "CustomStringField"; then
+  if checkExists "IwfWorkflowType" && checkExists "IwfGlobalWorkflowVersion" && checkExists "IwfExecutingStateIds" && checkExists "CustomKeywordField" && checkExists "CustomIntField" && checkExists "CustomBoolField" && checkExists "CustomDoubleField" && checkExists "CustomDatetimeField" && checkExists "CustomStringField" && checkExists "CustomKeywordArrayField"; then
     echo "All search attributes are registered"
     break
   fi
diff --git a/iwf/tests/test_persistence_data_attributes.py b/iwf/tests/test_persistence_data_attributes.py
@@ -63,7 +63,7 @@ def execute(
         return StateDecision.graceful_complete_workflow()
 
 
-class PersistenceWorkflow(ObjectWorkflow):
+class PersistenceDataAttributesWorkflow(ObjectWorkflow):
     def get_workflow_states(self) -> StateSchema:
         return StateSchema.with_starting_state(DataAttributeRWState())
 
@@ -85,14 +85,14 @@ def test_persistence_read(self, pers: Persistence):
         )
 
 
-class TestPersistence(unittest.TestCase):
+class TestPersistenceDataAttributes(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        wf = PersistenceWorkflow()
+        wf = PersistenceDataAttributesWorkflow()
         registry.add_workflow(wf)
         cls.client = Client(registry)
 
-    def test_persistence_workflow(self):
+    def test_persistence_data_attributes_workflow(self):
         wf_id = f"{inspect.currentframe().f_code.co_name}-{time.time_ns()}"
 
         start_options = WorkflowOptions(
@@ -102,10 +102,14 @@ def test_persistence_workflow(self):
             },
         )
 
-        self.client.start_workflow(PersistenceWorkflow, wf_id, 100, None, start_options)
-        self.client.get_simple_workflow_result_with_wait(wf_id, None)
+        self.client.start_workflow(
+            PersistenceDataAttributesWorkflow, wf_id, 100, None, start_options
+        )
+        self.client.wait_for_workflow_completion(wf_id, None)
 
-        res = self.client.invoke_rpc(wf_id, PersistenceWorkflow.test_persistence_read)
+        res = self.client.invoke_rpc(
+            wf_id, PersistenceDataAttributesWorkflow.test_persistence_read
+        )
         assert res == [
             final_initial_da_value_1,
             final_initial_da_value_2,
diff --git a/iwf/tests/test_persistence_search_attributes.py b/iwf/tests/test_persistence_search_attributes.py
diff --git a/iwf/unregistered_client.py b/iwf/unregistered_client.py
diff --git a/iwf/worker_service.py b/iwf/worker_service.py
