fix(versioning): fix exception getting feature states for edge identity post v2 versioning migration

api/edge_api/identities/models.py

@@ -18,6 +18,7 @@
 from environments.models import Environment
 from features.models import FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
+from features.versioning.versioning_service import get_environment_flags_dict
 from users.models import FFAdminUser
 from util.mappers import map_engine_identity_to_identity_document
 
@@ -80,40 +81,33 @@ def get_all_feature_states(
         )
         django_environment = Environment.objects.get(api_key=self.environment_api_key)
 
-        q = (
-            Q(version__isnull=False)
-            & Q(identity__isnull=True)
-            & (
-                Q(feature_segment__segment__id__in=segment_ids)
-                | Q(feature_segment__isnull=True)
+        feature_states: dict[str, FeatureState | FeatureStateModel] = (
+            get_environment_flags_dict(
+                environment=django_environment,
+                additional_filters=(
+                    Q(identity__isnull=True)
+                    & (
+                        Q(feature_segment__segment__id__in=segment_ids)
+                        | Q(feature_segment__isnull=True)
+                    )
+                ),
+                additional_select_related_args=[
+                    "feature",
+                    "feature_segment",
+                    "feature_segment__segment",
+                    "feature_state_value",
+                ],
+                additional_prefetch_related_args=[
+                    Prefetch(
+                        "multivariate_feature_state_values",
+                        queryset=MultivariateFeatureStateValue.objects.select_related(
+                            "multivariate_feature_option"
+                        ),
+                    )
+                ],
+                key_function=lambda fs: fs.feature.name,
             )
         )
-        environment_and_segment_feature_states = (
-            django_environment.feature_states.select_related(
-                "feature",
-                "feature_segment",
-                "feature_segment__segment",
-                "feature_state_value",
-            )
-            .prefetch_related(
-                Prefetch(
-                    "multivariate_feature_state_values",
-                    queryset=MultivariateFeatureStateValue.objects.select_related(
-                        "multivariate_feature_option"
-                    ),
-                )
-            )
-            .filter(q)
-        )
-
-        feature_states = {}
-        for feature_state in environment_and_segment_feature_states:
-            feature_name = feature_state.feature.name
-            if (
-                feature_name not in feature_states
-                or feature_state > feature_states[feature_name]
-            ):
-                feature_states[feature_name] = feature_state
 
         identity_feature_states = self.feature_overrides
         identity_feature_names = set()

api/features/versioning/versioning_service.py

@@ -14,9 +14,6 @@ def get_environment_flags_queryset(
     """
     Get a queryset of the latest live versions of an environments' feature states
     """
-    if environment.use_v2_feature_versioning:
-        return _get_feature_states_queryset(environment, feature_name)
-
     feature_states_list = get_environment_flags_list(environment, feature_name)
     return FeatureState.objects.filter(id__in=[fs.id for fs in feature_states_list])
 
@@ -39,6 +36,29 @@ def get_environment_flags_list(
     feature states. The logic to grab the latest version is then handled in python
     by building a dictionary. Returns a list of FeatureState objects.
     """
+    return list(
+        get_environment_flags_dict(
+            environment,
+            feature_name,
+            additional_filters,
+            additional_select_related_args,
+            additional_prefetch_related_args,
+        ).values()
+    )
+
+
+def get_environment_flags_dict(
+    environment: Environment,
+    feature_name: str = None,
+    additional_filters: Q = None,
+    additional_select_related_args: typing.Iterable[str] = None,
+    additional_prefetch_related_args: typing.Iterable[
+        typing.Union[str, Prefetch]
+    ] = None,
+    key_function: typing.Callable[[FeatureState], tuple] = None,
+) -> dict[tuple | str | int, FeatureState]:
+    key_function = key_function or _get_distinct_key
+
     feature_states = _get_feature_states_queryset(
         environment,
         feature_name,
@@ -47,24 +67,17 @@ def get_environment_flags_list(
         additional_prefetch_related_args,
     )
 
-    if environment.use_v2_feature_versioning:
-        return list(feature_states)
-
-    # Build up a dictionary in the form
-    # {(feature_id, feature_segment_id, identity_id): feature_state}
-    # and only keep the latest version for each feature.
+    # Build up a dictionary keyed off the relevant unique attributes as defined
+    # by the provided key function and only keep the highest priority feature state
+    # for each feature.
     feature_states_dict = {}
     for feature_state in feature_states:
-        key = (
-            feature_state.feature_id,
-            getattr(feature_state.feature_segment, "segment_id", None),
-            feature_state.identity_id,
-        )
+        key = key_function(feature_state)
         current_feature_state = feature_states_dict.get(key)
         if not current_feature_state or feature_state > current_feature_state:
             feature_states_dict[key] = feature_state
 
-    return list(feature_states_dict.values())
+    return feature_states_dict
 
 
 def get_current_live_environment_feature_version(
@@ -113,3 +126,13 @@ def _get_feature_states_queryset(
         queryset = queryset.filter(feature__name__iexact=feature_name)
 
     return queryset
+
+
+def _get_distinct_key(
+    feature_state: FeatureState,
+) -> tuple[int, int | None, int | None]:
+    return (
+        feature_state.feature_id,
+        getattr(feature_state.feature_segment, "segment_id", None),
+        feature_state.identity_id,
+    )

api/tests/unit/edge_api/identities/test_edge_identity_models.py

@@ -9,7 +9,9 @@
 from pytest_mock import MockerFixture
 
 from edge_api.identities.models import EdgeIdentity
-from features.models import FeatureSegment, FeatureState, FeatureStateValue
+from environments.models import Environment
+from features.models import Feature, FeatureSegment, FeatureState
+from features.versioning.tasks import enable_v2_versioning
 from features.workflows.core.models import ChangeRequest
 from segments.models import Segment
 
@@ -38,17 +40,10 @@ def test_get_all_feature_states_for_edge_identity_uses_segment_priorities(
     segment_override_p1 = FeatureState.objects.create(
         feature=feature, environment=environment, feature_segment=feature_segment_p1
     )
-    segment_override_p2 = FeatureState.objects.create(
+    FeatureState.objects.create(
         feature=feature, environment=environment, feature_segment=feature_segment_p2
     )
 
-    FeatureStateValue.objects.filter(feature_state=segment_override_p1).update(
-        string_value="p1"
-    )
-    FeatureStateValue.objects.filter(feature_state=segment_override_p2).update(
-        string_value="p2"
-    )
-
     identity_model = mocker.MagicMock(
         environment_api_key=environment.api_key, identity_features=[]
     )
@@ -479,3 +474,43 @@ def test_edge_identity_save_called_generate_audit_records_if_feature_override_up
             "identifier": edge_identity_model.identifier,
         }
     )
+
+
+def test_get_all_feature_states_post_v2_versioning_migration(
+    environment: Environment,
+    feature: Feature,
+    feature_state: FeatureState,
+    segment: Segment,
+    segment_featurestate: FeatureState,
+    edge_identity_model: EdgeIdentity,
+    mocker: MockerFixture,
+) -> None:
+    """
+    Specific test to reproduce an issue seen after migrating our staging environment to
+    v2 versioning.
+    """
+
+    # Given
+    # Multiple versions (old style versioning) of a given environment feature state and segment override
+    # to simulate the fact that we will be left with feature states that do NOT have a feature version
+    # (because only the latest versions get migrated to v2 versioning).
+    feature_state.clone(env=environment, live_from=timezone.now(), version=2)
+    v2_segment_override = segment_featurestate.clone(
+        env=environment, live_from=timezone.now(), version=2
+    )
+
+    enable_v2_versioning(environment.id)
+
+    edge_identity_dynamo_wrapper_mock = mocker.patch(
+        "edge_api.identities.models.EdgeIdentity.dynamo_wrapper",
+    )
+    edge_identity_dynamo_wrapper_mock.get_segment_ids.return_value = [segment.id]
+
+    # When
+    feature_states, identity_override_feature_names = (
+        edge_identity_model.get_all_feature_states()
+    )
+
+    # Then
+    assert len(feature_states) == 1
+    assert feature_states[0] == v2_segment_override