Merge branch 'main' into feat/add_api_usage_billing

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,41 +81,43 @@ 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)
-            )
+        # since identity overrides are included in the document retrieved from dynamo,
+        # we only want to retrieve the environment default and (relevant) segment overrides
+        # from the ORM.
+        additional_filters = Q(identity__isnull=True) & (
+            Q(feature_segment__segment__id__in=segment_ids)
+            | Q(feature_segment__isnull=True)
         )
-        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"
-                    ),
-                )
+
+        feature_states: dict[str, FeatureState | FeatureStateModel] = (
+            get_environment_flags_dict(
+                environment=django_environment,
+                additional_filters=additional_filters,
+                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"
+                        ),
+                    )
+                ],
+                # since we only want to retrieve the highest priority feature state,
+                # we key off the feature name instead of the default
+                # (feature_id, segment_id, identity_id). This will give us only e.g.
+                # the highest priority matching segment override for a given feature.
+                key_function=lambda fs: fs.feature.name,
             )
-            .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
-
+        # Since the identity overrides are the highest priority, we can now iterate
+        # over the dictionary and replace any feature states with those that have
+        # an identity override, stored against the identity in dynamo.
         identity_feature_states = self.feature_overrides
         identity_feature_names = set()
         for identity_feature_state in identity_feature_states:

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

@@ -6,10 +6,13 @@
 from django.utils import timezone
 from flag_engine.features.models import FeatureModel, FeatureStateModel
 from freezegun import freeze_time
+from pytest_django import DjangoAssertNumQueries
 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 +41,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 +475,45 @@ 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,
+    django_assert_num_queries: DjangoAssertNumQueries,
+) -> 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
+    with django_assert_num_queries(4):
+        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

api/tests/unit/features/versioning/test_unit_versioning_versioning_service.py

@@ -2,6 +2,7 @@
 
 from django.db.models import Q
 from django.utils import timezone
+from pytest_django import DjangoAssertNumQueries
 
 from environments.identities.models import Identity
 from environments.models import Environment
@@ -18,7 +19,9 @@
 
 
 def test_get_environment_flags_queryset_returns_only_latest_versions(
-    feature, environment
+    feature: Feature,
+    environment: Environment,
+    django_assert_num_queries: DjangoAssertNumQueries,
 ):
     # Given
     feature_state_v1 = FeatureState.objects.get(
@@ -31,7 +34,11 @@ def test_get_environment_flags_queryset_returns_only_latest_versions(
     feature_state_v1.clone(env=environment, as_draft=True)  # draft feature state
 
     # When
-    feature_states = get_environment_flags_queryset(environment=environment)
+    with django_assert_num_queries(2):
+        feature_states = get_environment_flags_queryset(environment=environment)
+
+        # trigger the queryset to execute and ensure the number of queries is correct
+        list(feature_states)
 
     # Then
     assert feature_states.count() == 1
@@ -115,6 +122,7 @@ def test_get_environment_flags_v2_versioning_returns_latest_live_versions_of_fea
     environment_v2_versioning: Environment,
     feature: Feature,
     admin_user: FFAdminUser,
+    django_assert_num_queries: DjangoAssertNumQueries,
 ) -> None:
     # Given
     # a second feature with its corresponding environment feature version
@@ -139,10 +147,11 @@ def test_get_environment_flags_v2_versioning_returns_latest_live_versions_of_fea
     environment_feature_1_version_2.publish(admin_user)
 
     # When
-    environment_feature_states = get_environment_flags_list(
-        environment=environment_v2_versioning,
-        additional_filters=Q(feature_segment=None, identity=None),
-    )
+    with django_assert_num_queries(2):
+        environment_feature_states = get_environment_flags_list(
+            environment=environment_v2_versioning,
+            additional_filters=Q(feature_segment=None, identity=None),
+        )
 
     # Then
     assert set(environment_feature_states) == {