feat: define python client search attribute

timl3136 · timl3136 · commit 7f01cf2e46d3 · 2026-02-04T21:28:59.000-08:00
Signed-off-by: Tim Li &lt;ltim@uber.com&gt;
diff --git a/cadence/__init__.py b/cadence/__init__.py
@@ -8,11 +8,23 @@
 from .client import Client
 from .worker import Registry
 from . import workflow
+from .search_attributes import (
+    SearchAttributeConverter,
+    SearchAttributeValue,
+    CADENCE_CHANGE_VERSION,
+    validate_search_attributes,
+)
 
 __version__ = "0.1.0"
 
 __all__ = [
+    # Core
     "Client",
     "Registry",
     "workflow",
+    # Search Attributes
+    "SearchAttributeConverter",
+    "SearchAttributeValue",
+    "CADENCE_CHANGE_VERSION",
+    "validate_search_attributes",
 ]
diff --git a/cadence/search_attributes.py b/cadence/search_attributes.py
@@ -0,0 +1,118 @@
+"""
+Search Attributes support for Cadence workflows.
+
+Search attributes are custom indexed fields attached to workflow executions
+that enable advanced visibility queries using SQL-like syntax.
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from typing import Any
+
+from cadence.api.v1.common_pb2 import Payload, SearchAttributes
+
+# Supported Python types for search attribute values (any JSON-serializable type)
+SearchAttributeValue = Any
+
+# CadenceChangeVersion is used as search attributes key to find workflows with specific change version.
+CADENCE_CHANGE_VERSION = "CadenceChangeVersion"
+
+class SearchAttributeConverter:
+    """
+    Converts between Python dictionaries and protobuf SearchAttributes.
+    """
+
+    def encode(self, attrs: dict[str, SearchAttributeValue]) -> SearchAttributes:
+        """
+        Encode a Python dictionary to a protobuf SearchAttributes message.
+
+        Args:
+            attrs: Dictionary mapping attribute names to Python values
+
+        Returns:
+            SearchAttributes protobuf message
+
+        Raises:
+            ValueError: If attrs is empty or contains unsupported types
+        """
+        if not attrs:
+            raise ValueError("search attributes is empty")
+
+        result = SearchAttributes()
+
+        for key, value in attrs.items():
+            try:
+                encoded_bytes = self._encode_value(value)
+                result.indexed_fields[key].CopyFrom(Payload(data=encoded_bytes))
+            except (TypeError, ValueError) as e:
+                raise ValueError(f"encode search attribute [{key}] error: {e}") from e
+
+        return result
+
+    def decode(
+        self,
+        attrs: SearchAttributes,
+        type_hints: dict[str, type] | None = None,
+    ) -> dict[str, Any]:
+        """
+        Decode a protobuf SearchAttributes message to a Python dictionary.
+
+        Args:
+            attrs: SearchAttributes protobuf message
+            type_hints: Optional type hints for datetime conversion
+
+        Returns:
+            Dictionary mapping attribute names to decoded values
+        """
+        if not attrs.indexed_fields:
+            return {}
+
+        type_hints = type_hints or {}
+        result: dict[str, Any] = {}
+
+        for key, payload in attrs.indexed_fields.items():
+            hint = type_hints.get(key)
+            result[key] = self._decode_value(payload.data, hint)
+
+        return result
+
+    def _encode_value(self, value: SearchAttributeValue) -> bytes:
+        """Encode a single value to JSON bytes."""
+        if isinstance(value, datetime):
+            return json.dumps(value.isoformat()).encode("utf-8")
+        else:
+            return json.dumps(value).encode("utf-8")
+
+    def _decode_value(self, data: bytes, type_hint: type | None = None) -> Any:
+        """Decode JSON bytes to a Python value."""
+        if not data:
+            return None
+
+        raw_value = json.loads(data.decode("utf-8"))
+
+        if type_hint is datetime and isinstance(raw_value, str):
+            return datetime.fromisoformat(raw_value)
+
+        return raw_value
+
+
+def validate_search_attributes(
+    attrs: dict[str, SearchAttributeValue],
+    allow_reserved_keys: bool = False,
+) -> None:
+    """
+    Validate search attributes before encoding.
+
+    Raises:
+        ValueError: If empty or uses reserved keys
+    """
+    if not attrs:
+        raise ValueError("search attributes is empty")
+
+    if not allow_reserved_keys and CADENCE_CHANGE_VERSION in attrs:
+        raise ValueError(
+            f"{CADENCE_CHANGE_VERSION} is a reserved key that cannot be set, "
+            "please use other key"
+        )
diff --git a/tests/cadence/test_search_attributes.py b/tests/cadence/test_search_attributes.py
@@ -0,0 +1,154 @@
+"""Tests for the search_attributes module."""
+
+import json
+from datetime import datetime, timezone
+
+import pytest
+
+from cadence.api.v1.common_pb2 import Payload, SearchAttributes
+from cadence.search_attributes import (
+    CADENCE_CHANGE_VERSION,
+    SearchAttributeConverter,
+    validate_search_attributes,
+)
+
+
+@pytest.fixture
+def converter() -> SearchAttributeConverter:
+    return SearchAttributeConverter()
+
+
+def make_search_attrs(attrs: dict[str, bytes]) -> SearchAttributes:
+    result = SearchAttributes()
+    for key, data in attrs.items():
+        result.indexed_fields[key].CopyFrom(Payload(data=data))
+    return result
+
+
+class TestEncode:
+    def test_string(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Status": "RUNNING"})
+        assert result.indexed_fields["Status"].data == b'"RUNNING"'
+
+    def test_int(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Count": 42})
+        assert result.indexed_fields["Count"].data == b"42"
+
+    def test_float(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Price": 99.99})
+        assert result.indexed_fields["Price"].data == b"99.99"
+
+    def test_bool(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"IsTrue": True, "IsFalse": False})
+        assert result.indexed_fields["IsTrue"].data == b"true"
+        assert result.indexed_fields["IsFalse"].data == b"false"
+
+    def test_datetime(self, converter: SearchAttributeConverter) -> None:
+        dt = datetime(2024, 1, 15, 10, 30, 0, tzinfo=timezone.utc)
+        result = converter.encode({"CreatedAt": dt})
+        assert json.loads(result.indexed_fields["CreatedAt"].data) == "2024-01-15T10:30:00+00:00"
+
+    def test_list_of_strings(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Tags": ["urgent", "customer"]})
+        assert json.loads(result.indexed_fields["Tags"].data) == ["urgent", "customer"]
+
+    def test_multiple_attributes(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Status": "RUNNING", "Count": 42})
+        assert len(result.indexed_fields) == 2
+
+    def test_empty_raises(self, converter: SearchAttributeConverter) -> None:
+        with pytest.raises(ValueError, match="search attributes is empty"):
+            converter.encode({})
+
+    def test_list_of_ints(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Counts": [1, 2, 3]})
+        assert json.loads(result.indexed_fields["Counts"].data) == [1, 2, 3]
+
+    def test_nested_dict(self, converter: SearchAttributeConverter) -> None:
+        result = converter.encode({"Meta": {"key": "value"}})
+        assert json.loads(result.indexed_fields["Meta"].data) == {"key": "value"}
+
+
+class TestDecode:
+    def test_string(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"Status": b'"RUNNING"'})
+        assert converter.decode(proto) == {"Status": "RUNNING"}
+
+    def test_int(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"Count": b"42"})
+        assert converter.decode(proto) == {"Count": 42}
+
+    def test_float(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"Price": b"99.99"})
+        assert converter.decode(proto) == {"Price": 99.99}
+
+    def test_bool(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"IsTrue": b"true", "IsFalse": b"false"})
+        assert converter.decode(proto) == {"IsTrue": True, "IsFalse": False}
+
+    def test_list(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"Tags": b'["a", "b"]'})
+        assert converter.decode(proto) == {"Tags": ["a", "b"]}
+
+    def test_datetime_without_hint(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"CreatedAt": b'"2024-01-15T10:30:00"'})
+        assert converter.decode(proto) == {"CreatedAt": "2024-01-15T10:30:00"}
+
+    def test_datetime_with_hint(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"CreatedAt": b'"2024-01-15T10:30:00+00:00"'})
+        result = converter.decode(proto, type_hints={"CreatedAt": datetime})
+        assert result["CreatedAt"] == datetime(2024, 1, 15, 10, 30, 0, tzinfo=timezone.utc)
+
+    def test_empty_search_attributes(self, converter: SearchAttributeConverter) -> None:
+        assert converter.decode(SearchAttributes()) == {}
+
+    def test_empty_payload(self, converter: SearchAttributeConverter) -> None:
+        proto = make_search_attrs({"Empty": b""})
+        assert converter.decode(proto) == {"Empty": None}
+
+
+class TestRoundTrip:
+    @pytest.mark.parametrize(
+        "attrs,type_hints",
+        [
+            ({"Status": "RUNNING"}, None),
+            ({"Count": 42}, None),
+            ({"Price": 99.99}, None),
+            ({"IsPriority": True}, None),
+            ({"Tags": ["a", "b"]}, None),
+            ({"Counts": [1, 2, 3]}, None),
+            (
+                {"CreatedAt": datetime(2024, 1, 15, 10, 30, 0, tzinfo=timezone.utc)},
+                {"CreatedAt": datetime},
+            ),
+        ],
+    )
+    def test_roundtrip(
+        self,
+        converter: SearchAttributeConverter,
+        attrs: dict[str, type],
+        type_hints: dict[str, type] | None,
+    ) -> None:
+        encoded = converter.encode(attrs)
+        decoded = converter.decode(encoded, type_hints=type_hints)
+        assert decoded == attrs
+
+
+class TestValidate:
+    def test_valid(self) -> None:
+        validate_search_attributes({"Status": "RUNNING", "Count": 42})
+
+    def test_empty_raises(self) -> None:
+        with pytest.raises(ValueError, match="search attributes is empty"):
+            validate_search_attributes({})
+
+    def test_reserved_key_raises(self) -> None:
+        with pytest.raises(ValueError, match="reserved key"):
+            validate_search_attributes({CADENCE_CHANGE_VERSION: ["v1"]})
+
+    def test_reserved_key_allowed(self) -> None:
+        validate_search_attributes({CADENCE_CHANGE_VERSION: ["v1"]}, allow_reserved_keys=True)
+
+
+def test_cadence_change_version_constant() -> None:
+    assert CADENCE_CHANGE_VERSION == "CadenceChangeVersion"
