From 5392480c3e35fd689347a80714901d4f70116367 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Mon, 27 Nov 2023 16:31:43 +0000
Subject: [PATCH] fix: resolve environment N+1 caused by feature versioning v2
 (#3040)

---
 api/environments/identities/models.py         |  1 +
 .../identities/tests/test_views.py            | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/api/environments/identities/models.py b/api/environments/identities/models.py
index be681e0e76c3..61aca21719d6 100644
--- a/api/environments/identities/models.py
+++ b/api/environments/identities/models.py
@@ -97,6 +97,7 @@ def get_all_feature_states(
             full_query &= additional_filters
 
         select_related_args = [
+            "environment",
             "feature",
             "feature_state_value",
             "feature_segment",
diff --git a/api/environments/identities/tests/test_views.py b/api/environments/identities/tests/test_views.py
index 6086c707b002..b9552d671391 100644
--- a/api/environments/identities/tests/test_views.py
+++ b/api/environments/identities/tests/test_views.py
@@ -7,6 +7,7 @@
 from core.constants import FLAGSMITH_UPDATED_AT_HEADER
 from django.test import override_settings
 from django.urls import reverse
+from django.utils import timezone
 from rest_framework import status
 from rest_framework.test import APIClient, APITestCase
 
@@ -849,6 +850,35 @@ def test_get_identities_request_includes_updated_at_header(self):
             self.environment.updated_at.timestamp()
         )
 
+    def test_get_identities_nplus1(self) -> None:
+        """
+        Specific test to reproduce N+1 issue found after deployment of
+        v2 feature versioning.
+
+        TODO: move this (and other tests) to /api/tests/...
+        """
+        url = "%s?identifier=%s" % (
+            reverse("api-v1:sdk-identities"),
+            self.identity.identifier,
+        )
+
+        # let's get a baseline with only a single version of a flag
+        with self.assertNumQueries(6):
+            self.client.get(url)
+
+        # now let's create some new versions of the same flag
+        # and verify that the query count doesn't increase
+        v1_flag = FeatureState.objects.get(
+            environment=self.environment, feature=self.feature_1
+        )
+        now = timezone.now()
+        for i in range(2, 13):
+            v1_flag.clone(env=self.environment, version=i, live_from=now)
+
+        # in fact it is lower because the
+        with self.assertNumQueries(5):
+            self.client.get(url)
+
 
 def test_get_identities_with_hide_sensitive_data_with_feature_name(
     environment, feature, identity, api_client