feat: Support feature metadata (#163)

Author: khvn26

flagsmith/mappers.py

@@ -22,14 +22,15 @@
 )
 from flagsmith.models import Segment
 from flagsmith.types import (
+    FeatureMetadata,
     SDKEvaluationContext,
     SegmentMetadata,
     StreamEvent,
     TraitConfig,
 )
 
 OverrideKey = typing.Tuple[
-    str,
+    int,
     str,
     bool,
     typing.Any,
@@ -148,7 +149,7 @@ def map_environment_document_to_context(
 
 def _map_identity_overrides_to_segments(
     identity_overrides: list[IdentityModel],
-) -> dict[str, SegmentContext[SegmentMetadata]]:
+) -> dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]:
     features_to_identifiers: typing.Dict[
         OverridesKey,
         typing.List[str],
@@ -159,7 +160,7 @@ def _map_identity_overrides_to_segments(
             continue
         overrides_key = tuple(
             (
-                str(feature_state["feature"]["id"]),
+                feature_state["feature"]["id"],
                 feature_state["feature"]["name"],
                 feature_state["enabled"],
                 feature_state["feature_state_value"],
@@ -170,7 +171,13 @@ def _map_identity_overrides_to_segments(
             )
         )
         features_to_identifiers[overrides_key].append(identity_override["identifier"])
-    segment_contexts: typing.Dict[str, SegmentContext[SegmentMetadata]] = {}
+    segment_contexts: typing.Dict[
+        str,
+        SegmentContext[
+            SegmentMetadata,
+            FeatureMetadata,
+        ],
+    ] = {}
     for overrides_key, identifiers in features_to_identifiers.items():
         # Create a segment context for each unique set of overrides
         # Generate a unique key to avoid collisions
@@ -193,13 +200,14 @@ def _map_identity_overrides_to_segments(
             overrides=[
                 {
                     "key": "",  # Identity overrides never carry multivariate options
-                    "feature_key": feature_key,
+                    "feature_key": str(flagsmith_id),
                     "name": feature_name,
                     "enabled": feature_enabled,
                     "value": feature_value,
                     "priority": float("-inf"),  # Highest possible priority
+                    "metadata": {"flagsmith_id": flagsmith_id},
                 }
-                for feature_key, feature_name, feature_enabled, feature_value in overrides_key
+                for flagsmith_id, feature_name, feature_enabled, feature_value in overrides_key
             ],
             metadata=SegmentMetadata(source="identity_overrides"),
         )
@@ -230,16 +238,18 @@ def _map_environment_document_rules_to_context_rules(
 
 def _map_environment_document_feature_states_to_feature_contexts(
     feature_states: list[FeatureStateModel],
-) -> typing.Iterable[FeatureContext]:
+) -> typing.Iterable[FeatureContext[FeatureMetadata]]:
     for feature_state in feature_states:
-        feature_context = FeatureContext(
+        metadata: FeatureMetadata = {"flagsmith_id": feature_state["feature"]["id"]}
+        feature_context = FeatureContext[FeatureMetadata](
             key=str(
                 feature_state.get("django_id") or feature_state["featurestate_uuid"]
             ),
             feature_key=str(feature_state["feature"]["id"]),
             name=feature_state["feature"]["name"],
             enabled=feature_state["enabled"],
             value=feature_state["feature_state_value"],
+            metadata=metadata,
         )
         if multivariate_feature_state_values := feature_state.get(
             "multivariate_feature_state_values"

flagsmith/models.py

@@ -3,11 +3,9 @@
 import typing
 from dataclasses import dataclass, field
 
-from flag_engine.result.types import FlagResult
-
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithFeatureDoesNotExistError
-from flagsmith.types import SDKEvaluationResult
+from flagsmith.types import SDKEvaluationResult, SDKFlagResult
 
 
 @dataclass
@@ -30,13 +28,18 @@ class Flag(BaseFlag):
     @classmethod
     def from_evaluation_result(
         cls,
-        flag: FlagResult,
+        flag_result: SDKFlagResult,
     ) -> Flag:
-        return Flag(
-            enabled=flag["enabled"],
-            value=flag["value"],
-            feature_name=flag["name"],
-            feature_id=int(flag["feature_key"]),
+        if metadata := flag_result.get("metadata"):
+            return Flag(
+                enabled=flag_result["enabled"],
+                value=flag_result["value"],
+                feature_name=flag_result["name"],
+                feature_id=metadata["flagsmith_id"],
+            )
+        raise ValueError(
+            "FlagResult metadata is missing. Cannot create Flag instance. "
+            "This means a bug in the SDK, please report it."
         )
 
     @classmethod
@@ -64,13 +67,9 @@ def from_evaluation_result(
     ) -> Flags:
         return cls(
             flags={
-                flag_name: Flag(
-                    enabled=flag["enabled"],
-                    value=flag["value"],
-                    feature_name=flag["name"],
-                    feature_id=int(flag["feature_key"]),
-                )
-                for flag_name, flag in evaluation_result["flags"].items()
+                flag_name: flag
+                for flag_name, flag_result in evaluation_result["flags"].items()
+                if (flag := Flag.from_evaluation_result(flag_result))
             },
             default_flag_handler=default_flag_handler,
             _analytics_processor=analytics_processor,

flagsmith/types.py

@@ -3,7 +3,7 @@
 
 from flag_engine.context.types import EvaluationContext
 from flag_engine.engine import ContextValue
-from flag_engine.result.types import EvaluationResult
+from flag_engine.result.types import EvaluationResult, FlagResult
 from typing_extensions import NotRequired, TypeAlias
 
 _JsonScalarType: TypeAlias = typing.Union[
@@ -44,5 +44,11 @@ class SegmentMetadata(typing.TypedDict):
     """The source of the segment, e.g. 'api', 'identity_overrides'."""
 
 
-SDKEvaluationContext = EvaluationContext[SegmentMetadata]
-SDKEvaluationResult = EvaluationResult[SegmentMetadata]
+class FeatureMetadata(typing.TypedDict):
+    flagsmith_id: int
+    """The ID of the feature used in Flagsmith API."""
+
+
+SDKEvaluationContext = EvaluationContext[SegmentMetadata, FeatureMetadata]
+SDKEvaluationResult = EvaluationResult[SegmentMetadata, FeatureMetadata]
+SDKFlagResult = FlagResult[FeatureMetadata]

poetry.lock

@@ -10,7 +10,7 @@ documentation = "https://docs.flagsmith.com"
 packages = [{ include = "flagsmith" }]
 
 [tool.poetry.dependencies]
-flagsmith-flag-engine = "^9.0.0"
+flagsmith-flag-engine = "^9.1.0"
 python = ">=3.9,<4"
 requests = "^2.32.3"
 requests-futures = "^1.0.1"

pyproject.toml

@@ -166,6 +166,7 @@ def test_get_identity_flags_uses_local_environment_when_available(
                 "enabled": True,
                 "value": "some-feature-state-value",
                 "feature_key": "1",
+                "metadata": {"flagsmith_id": 1},
             }
         },
         "segments": [],

tests/test_flagsmith.py

@@ -1,24 +1,24 @@
 import typing
 
 import pytest
-from flag_engine.result.types import FlagResult
 
 from flagsmith.models import Flag, Flags
-from flagsmith.types import SDKEvaluationResult
+from flagsmith.types import SDKEvaluationResult, SDKFlagResult
 
 
 def test_flag_from_evaluation_result() -> None:
     # Given
-    flag_result: FlagResult = {
+    flag_result: SDKFlagResult = {
         "enabled": True,
         "feature_key": "123",
         "name": "test_feature",
         "reason": "DEFAULT",
         "value": "test-value",
+        "metadata": {"flagsmith_id": 123},
     }
 
     # When
-    flag: Flag = Flag.from_evaluation_result(flag_result)
+    flag = Flag.from_evaluation_result(flag_result)
 
     # Then
     assert flag.enabled is True
@@ -29,9 +29,9 @@ def test_flag_from_evaluation_result() -> None:
 
 
 @pytest.mark.parametrize(
-    "flags_result,expected_count,expected_names",
+    "flags_result,expected_names",
     [
-        ({}, 0, []),
+        ({}, []),
         (
             {
                 "feature1": {
@@ -40,9 +40,9 @@ def test_flag_from_evaluation_result() -> None:
                     "name": "feature1",
                     "reason": "DEFAULT",
                     "value": "value1",
+                    "metadata": {"flagsmith_id": 1},
                 }
             },
-            1,
             ["feature1"],
         ),
         (
@@ -53,9 +53,9 @@ def test_flag_from_evaluation_result() -> None:
                     "name": "feature1",
                     "reason": "DEFAULT",
                     "value": "value1",
+                    "metadata": {"flagsmith_id": 1},
                 }
             },
-            1,
             ["feature1"],
         ),
         (
@@ -66,30 +66,31 @@ def test_flag_from_evaluation_result() -> None:
                     "name": "feature1",
                     "reason": "DEFAULT",
                     "value": "value1",
+                    "metadata": {"flagsmith_id": 1},
                 },
                 "feature2": {
                     "enabled": True,
                     "feature_key": "2",
                     "name": "feature2",
                     "reason": "DEFAULT",
                     "value": "value2",
+                    "metadata": {"flagsmith_id": 2},
                 },
                 "feature3": {
                     "enabled": True,
                     "feature_key": "3",
                     "name": "feature3",
                     "reason": "DEFAULT",
                     "value": 42,
+                    "metadata": {"flagsmith_id": 3},
                 },
             },
-            3,
             ["feature1", "feature2", "feature3"],
         ),
     ],
 )
 def test_flags_from_evaluation_result(
-    flags_result: typing.Dict[str, FlagResult],
-    expected_count: int,
+    flags_result: typing.Dict[str, SDKFlagResult],
     expected_names: typing.List[str],
 ) -> None:
     # Given
@@ -106,13 +107,10 @@ def test_flags_from_evaluation_result(
     )
 
     # Then
-    assert len(flags.flags) == expected_count
-
-    for name in expected_names:
-        assert name in flags.flags
-        flag: Flag = flags.flags[name]
-        assert isinstance(flag, Flag)
-        assert flag.feature_name == name
+    assert set(flags.flags.keys()) == set(expected_names)
+    assert set(flag.feature_name for flag in flags.flags.values()) == set(
+        expected_names
+    )
 
 
 @pytest.mark.parametrize(
@@ -130,16 +128,32 @@ def test_flag_from_evaluation_result_value_types(
     value: typing.Any, expected: typing.Any
 ) -> None:
     # Given
-    flag_result: FlagResult = {
+    flag_result: SDKFlagResult = {
         "enabled": True,
         "feature_key": "123",
         "name": "test_feature",
         "reason": "DEFAULT",
         "value": value,
+        "metadata": {"flagsmith_id": 123},
     }
 
     # When
-    flag: Flag = Flag.from_evaluation_result(flag_result)
+    flag = Flag.from_evaluation_result(flag_result)
 
     # Then
     assert flag.value == expected
+
+
+def test_flag_from_evaluation_result_missing_metadata__raises_expected() -> None:
+    # Given
+    flag_result: SDKFlagResult = {
+        "enabled": True,
+        "feature_key": "123",
+        "name": "test_feature",
+        "reason": "DEFAULT",
+        "value": "test-value",
+    }
+
+    # When & Then
+    with pytest.raises(ValueError):
+        Flag.from_evaluation_result(flag_result)