fix: N+1 on segment overrides for environment-document endpoint (#3512)

api/environments/managers.py

@@ -46,7 +46,7 @@ def filter_for_document_builder(self, *args, **kwargs):
                 Prefetch(
                     "project__segments__feature_segments__feature_states",
                     queryset=FeatureState.objects.select_related(
-                        "feature", "feature_state_value"
+                        "feature", "feature_state_value", "environment"
                     ),
                 ),
                 Prefetch(

api/tests/unit/environments/test_unit_environments_views_sdk_environment.py

@@ -5,7 +5,9 @@
 from rest_framework.test import APIClient
 
 from environments.models import Environment, EnvironmentAPIKey
-from features.models import Feature
+from features.feature_types import MULTIVARIATE
+from features.models import Feature, FeatureSegment, FeatureState
+from features.multivariate.models import MultivariateFeatureOption
 from segments.models import Condition, Segment, SegmentRule
 
 
@@ -22,7 +24,7 @@ def test_get_environment_document(
     client.credentials(HTTP_X_ENVIRONMENT_KEY=api_key.key)
 
     # and some other sample data to make sure we're testing all of the document
-    Feature.objects.create(name="test_feature", project=project)
+    feature = Feature.objects.create(name="test_feature", project=project)
     for i in range(10):
         segment = Segment.objects.create(project=project)
         segment_rule = SegmentRule.objects.create(
@@ -44,11 +46,37 @@ def test_get_environment_document(
             rule=nested_rule,
         )
 
+        # Let's create segment override for each segment too
+        feature_segment = FeatureSegment.objects.create(
+            segment=segment, feature=feature, environment=environment
+        )
+        FeatureState.objects.create(
+            feature=feature,
+            environment=environment,
+            feature_segment=feature_segment,
+            enabled=True,
+        )
+
+    for i in range(10):
+        mv_feature = Feature.objects.create(
+            name=f"mv_feature_{i}", project=project, type=MULTIVARIATE
+        )
+        MultivariateFeatureOption.objects.create(
+            feature=mv_feature,
+            default_percentage_allocation=10,
+            string_value="option-1",
+        )
+        MultivariateFeatureOption.objects.create(
+            feature=mv_feature,
+            default_percentage_allocation=10,
+            string_value="option-2",
+        )
+
     # and the relevant URL to get an environment document
     url = reverse("api-v1:environment-document")
 
     # When
-    with django_assert_num_queries(11):
+    with django_assert_num_queries(13):
         response = client.get(url)
 
     # Then

api/util/mappers/engine.py

@@ -413,6 +413,9 @@ def _get_prioritised_feature_states(
 ) -> List["FeatureState"]:
     prioritised_feature_state_by_feature_id = {}
     for feature_state in feature_states:
+        # TODO: this call to is_live was causing an N+1 issue.
+        #  For now, we have solved it with an extra select_related, but
+        #  there is probably a neater solution here.
         if not feature_state.is_live:
             continue
         if existing_feature_state := prioritised_feature_state_by_feature_id.get(