IWF-607: Implement Search Attributes store and API (#83)

* IWF-607: Refactor Persistence

* IWF-607: Lint

* IWF-607: Fix types

* IWF-607: Fix

* IWF-607: Create client methods

Author: lwolczynski

iwf-idl

@@ -5,10 +5,17 @@
 
 from iwf.client_options import ClientOptions
 from iwf.errors import InvalidArgumentError
+from iwf.iwf_api.models import (
+    SearchAttribute,
+    SearchAttributeKeyAndType,
+    SearchAttributeValueType,
+)
+from iwf.iwf_api.types import Unset
 from iwf.registry import Registry
 from iwf.reset_workflow_type_and_options import ResetWorkflowTypeAndOptions
 from iwf.stop_workflow_options import StopWorkflowOptions
 from iwf.unregistered_client import UnregisteredClient, UnregisteredWorkflowOptions
+from iwf.utils.iwf_typing import unset_to_none
 from iwf.workflow import ObjectWorkflow, get_workflow_type_by_class
 from iwf.workflow_options import WorkflowOptions
 from iwf.workflow_state import (
@@ -211,3 +218,106 @@ def skip_timer_at_command_index(
             state_execution_number,
             timer_command_index,
         )
+
+    def get_all_search_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        workflow_run_id: Optional[str] = None,
+    ):
+        return self._do_get_workflow_search_attributes(
+            workflow_class, workflow_id, workflow_run_id
+        )
+
+    def get_workflow_search_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        attribute_keys: list[str],
+        workflow_run_id: Optional[str] = None,
+    ):
+        if not attribute_keys:
+            raise ValueError(
+                "attribute_keys must contain at least one entry, or use get_all_search_attributes API to get all"
+            )
+        return self._do_get_workflow_search_attributes(
+            workflow_class, workflow_id, workflow_run_id, attribute_keys
+        )
+
+    def _do_get_workflow_search_attributes(
+        self,
+        workflow_class: type[ObjectWorkflow],
+        workflow_id: str,
+        workflow_run_id: Optional[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:
+            non_existing_search_attribute_list: list[str] = []
+            for attribute_key in attribute_keys:
+                if attribute_key not in search_attribute_types:
+                    non_existing_search_attribute_list.append(attribute_key)
+
+            if non_existing_search_attribute_list:
+                raise InvalidArgumentError(
+                    f"Search attributes not registered: {','.join(non_existing_search_attribute_list)}"
+                )
+
+        key_and_types: list[SearchAttributeKeyAndType] = []
+        if attribute_keys is None:
+            for attribute_key, sa_type in search_attribute_types.items():
+                key_and_types.append(SearchAttributeKeyAndType(attribute_key, sa_type))
+        else:
+            for attribute_key in attribute_keys:
+                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
+        )
+
+        response_sas = response.search_attributes
+
+        # TODO: troubleshoot why unset_to_none doesn't work as expected with lists
+        if isinstance(response_sas, Unset) or response_sas is None:
+            raise RuntimeError("search attributes not returned")
+
+        result: dict[str, Any] = {}
+
+        for response_sa in response_sas:
+            response_sa_key = unset_to_none(response_sa.key)
+            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)
+            result[response_sa_key] = value
+
+        return result
+
+    @staticmethod
+    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}")

iwf/client.py

@@ -0,0 +1,51 @@
+from typing import Any, Optional, Union
+
+from iwf.errors import WorkflowDefinitionError
+from iwf.iwf_api.models import EncodedObject
+from iwf.object_encoder import ObjectEncoder
+
+
+class DataAttributes:
+    _type_store: dict[str, Optional[type]]
+    _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]],
+        object_encoder: ObjectEncoder,
+        current_values: dict[str, Union[EncodedObject, None]],
+    ):
+        self._object_encoder = object_encoder
+        self._type_store = type_store
+        self._current_values = current_values
+        self._updated_values_to_return = {}
+
+    def get_data_attribute(self, key: str) -> Any:
+        if key not in self._type_store:
+            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[key]
+        return self._object_encoder.decode(encoded_object, registered_type)
+
+    def set_data_attribute(self, key: str, value: Any):
+        if key not in self._type_store:
+            raise WorkflowDefinitionError(f"data attribute %s is not registered {key}")
+
+        registered_type = self._type_store[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}"
+            )
+
+        encoded_value = self._object_encoder.encode(value)
+        self._current_values[key] = encoded_value
+        self._updated_values_to_return[key] = encoded_value
+
+    def get_updated_values_to_return(self) -> dict[str, EncodedObject]:
+        return self._updated_values_to_return

iwf/data_attributes.py

@@ -1,51 +1,65 @@
-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.data_attributes import DataAttributes
+from iwf.search_attributes import SearchAttributes
 
 
 class Persistence:
-    _type_store: dict[str, Optional[type]]
-    _object_encoder: ObjectEncoder
-    _current_values: dict[str, Union[EncodedObject, None]]
-    _updated_values_to_return: dict[str, EncodedObject]
+    _data_attributes: DataAttributes
+    _search_attributes: SearchAttributes
 
     def __init__(
         self,
-        type_store: dict[str, Optional[type]],
-        object_encoder: ObjectEncoder,
-        current_values: dict[str, Union[EncodedObject, None]],
+        data_attributes: DataAttributes,
+        search_attributes: SearchAttributes,
     ):
-        self._object_encoder = object_encoder
-        self._type_store = type_store
-        self._current_values = current_values
-        self._updated_values_to_return = {}
+        self._data_attributes = data_attributes
+        self._search_attributes = search_attributes
 
     def get_data_attribute(self, key: str) -> Any:
-        if key not in self._type_store:
-            raise WorkflowDefinitionError(f"data attribute %s is not registered {key}")
+        return self._data_attributes.get_data_attribute(key)
 
-        encoded_object = self._current_values.get(key)
-        if encoded_object is None:
-            return None
+    def set_data_attribute(self, key: str, value: Any):
+        self._data_attributes.set_data_attribute(key, value)
 
-        registered_type = self._type_store[key]
-        return self._object_encoder.decode(encoded_object, registered_type)
+    def get_search_attribute_int64(self, key: str) -> Union[None, int]:
+        return self._search_attributes.get_search_attribute_int64(key)
 
-    def set_data_attribute(self, key: str, value: Any):
-        if key not in self._type_store:
-            raise WorkflowDefinitionError(f"data attribute %s is not registered {key}")
+    def set_search_attribute_int64(self, key: str, value: 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):
+        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):
+        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):
+        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):
+        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)
 
-        registered_type = self._type_store[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}"
-            )
+    def set_search_attribute_datetime(self, key: str, value: str):
+        self._search_attributes.set_search_attribute_datetime(key, value)
 
-        encoded_value = self._object_encoder.encode(value)
-        self._current_values[key] = encoded_value
-        self._updated_values_to_return[key] = encoded_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 get_updated_values_to_return(self) -> dict[str, EncodedObject]:
-        return self._updated_values_to_return
+    def set_search_attribute_keyword_array(self, key: str, value: list[str]):
+        self._search_attributes.set_search_attribute_keyword_array(key, value)

iwf/persistence.py

@@ -2,26 +2,30 @@
 from enum import Enum
 from typing import List, Optional
 
+from iwf.iwf_api.models import SearchAttributeValueType
+
 
 class PersistenceFieldType(Enum):
     DataAttribute = 1
-    # SearchAttribute = 2
+    SearchAttribute = 2
 
 
 @dataclass
 class PersistenceField:
     key: str
     field_type: PersistenceFieldType
     value_type: Optional[type]
-    # search_attribute_type: Optional[SearchAttributeValueType] = None
+    search_attribute_type: Optional[SearchAttributeValueType] = None
 
     @classmethod
     def data_attribute_def(cls, key: str, value_type: Optional[type]):
         return PersistenceField(key, PersistenceFieldType.DataAttribute, value_type)
 
-    # @classmethod
-    # def search_attribute_def(cls, key: str, sa_type: SearchAttributeValueType):
-    #     return PersistenceField(key, PersistenceFieldType.SearchAttribute, sa_type)
+    @classmethod
+    def search_attribute_def(cls, key: str, sa_type: SearchAttributeValueType):
+        return PersistenceField(
+            key, PersistenceFieldType.SearchAttribute, None, sa_type
+        )
 
 
 @dataclass

iwf/persistence_schema.py

@@ -2,6 +2,7 @@
 
 from iwf.communication_schema import CommunicationMethodType
 from iwf.errors import InvalidArgumentError, WorkflowDefinitionError
+from iwf.iwf_api.models import SearchAttributeValueType
 from iwf.persistence_schema import PersistenceFieldType
 from iwf.rpc import RPCInfo
 from iwf.type_store import TypeStore, Type
@@ -16,6 +17,7 @@ class Registry:
     _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]]]
+    _search_attribute_types: dict[str, dict[str, SearchAttributeValueType]]
     _rpc_infos: dict[str, dict[str, RPCInfo]]
 
     def __init__(self):
@@ -25,6 +27,7 @@ def __init__(self):
         self._internal_channel_type_store = dict()
         self._signal_channel_type_store = dict()
         self._data_attribute_types = dict()
+        self._search_attribute_types = {}
         self._rpc_infos = dict()
 
     def add_workflow(self, wf: ObjectWorkflow):
@@ -33,6 +36,7 @@ def add_workflow(self, wf: ObjectWorkflow):
         self._register_internal_channels(wf)
         self._register_signal_channels(wf)
         self._register_data_attributes(wf)
+        self._register_search_attributes(wf)
         self._register_workflow_rpcs(wf)
 
     def add_workflows(self, *wfs: ObjectWorkflow):
@@ -73,6 +77,11 @@ def get_signal_channel_types(self, wf_type: str) -> dict[str, Optional[type]]:
     def get_data_attribute_types(self, wf_type: str) -> dict[str, Optional[type]]:
         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_rpc_infos(self, wf_type: str) -> dict[str, RPCInfo]:
         return self._rpc_infos[wf_type]
 
@@ -112,6 +121,23 @@ def _register_data_attributes(self, wf: ObjectWorkflow):
                 types[field.key] = field.value_type
         self._data_attribute_types[wf_type] = types
 
+    def _register_search_attributes(self, wf: ObjectWorkflow):
+        wf_type = get_workflow_type(wf)
+        types: dict[str, SearchAttributeValueType] = {}
+        for field in wf.get_persistence_schema().persistence_fields:
+            if field.field_type == PersistenceFieldType.SearchAttribute:
+                sa_type = field.search_attribute_type
+                if sa_type is None:
+                    raise WorkflowDefinitionError(
+                        f"Found search attribute {field.key} with no type set"
+                    )
+                if field.key in types:
+                    raise WorkflowDefinitionError(
+                        f"Search attribute {field.key} already exists"
+                    )
+                types[field.key] = sa_type
+        self._search_attribute_types[wf_type] = types
+
     def _register_workflow_state(self, wf):
         wf_type = get_workflow_type(wf)
         state_map = {}