fix(versioning): segment overrides limit (#4007)

api/environments/models.py

@@ -181,7 +181,7 @@ def clone(self, name: str, api_key: str = None) -> "Environment":
             # Grab the latest feature versions from the source environment.
             latest_environment_feature_versions = (
                 EnvironmentFeatureVersion.objects.get_latest_versions_as_queryset(
-                    environment=self
+                    environment_id=self.id
                 )
             )
 

api/environments/views.py

@@ -1,6 +1,6 @@
 import logging
 
-from django.db.models import Count
+from django.db.models import Count, Q
 from django.utils.decorators import method_decorator
 from drf_yasg import openapi
 from drf_yasg.utils import no_body, swagger_auto_schema
@@ -17,6 +17,7 @@
     NestedEnvironmentPermissions,
 )
 from environments.sdk.schemas import SDKEnvironmentDocumentModel
+from features.versioning.models import EnvironmentFeatureVersion
 from features.versioning.tasks import (
     disable_v2_versioning,
     enable_v2_versioning,
@@ -108,9 +109,28 @@ def get_queryset(self):
         queryset = Environment.objects.all()
 
         if self.action == "retrieve":
-            queryset = queryset.annotate(
-                total_segment_overrides=Count("feature_segments")
+            # Since we don't have the environment at this stage, we would need to query the database
+            # regardless, so it seems worthwhile to just query the database for the latest versions
+            # and use their existence as a proxy to whether v2 feature versioning is enabled.
+            latest_versions = EnvironmentFeatureVersion.objects.get_latest_versions_by_environment_api_key(
+                environment_api_key=self.kwargs["api_key"]
             )
+            if latest_versions:
+                # if there are latest versions (and hence v2 feature versioning is enabled), then
+                # we need to ensure that we're only counting the feature segments for those
+                # latest versions against the limits.
+                queryset = queryset.annotate(
+                    total_segment_overrides=Count(
+                        "feature_segments",
+                        filter=Q(
+                            feature_segments__environment_feature_version__in=latest_versions
+                        ),
+                    )
+                )
+            else:
+                queryset = queryset.annotate(
+                    total_segment_overrides=Count("feature_segments")
+                )
 
         return queryset
 

api/features/managers.py

@@ -33,8 +33,10 @@ def get_live_feature_states(
 
         qs_filter = Q(environment=environment, deleted_at__isnull=True)
         if environment.use_v2_feature_versioning:
-            latest_versions = EnvironmentFeatureVersion.objects.get_latest_versions(
-                environment
+            latest_versions = (
+                EnvironmentFeatureVersion.objects.get_latest_versions_by_environment_id(
+                    environment.id
+                )
             )
             latest_version_uuids = [efv.uuid for efv in latest_versions]
 

api/features/versioning/managers.py

@@ -5,7 +5,6 @@
 from softdelete.models import SoftDeleteManager
 
 if typing.TYPE_CHECKING:
-    from environments.models import Environment
     from features.versioning.models import EnvironmentFeatureVersion
 
 
@@ -14,16 +13,22 @@
 
 
 class EnvironmentFeatureVersionManager(SoftDeleteManager):
-    def get_latest_versions(self, environment: "Environment") -> RawQuerySet:
+    def get_latest_versions_by_environment_id(self, environment_id: int) -> RawQuerySet:
         """
         Get the latest EnvironmentFeatureVersion objects for a given environment.
         """
-        return self.raw(
-            get_latest_versions_sql, params={"environment_id": environment.id}
-        )
+        return self._get_latest_versions(environment_id=environment_id)
+
+    def get_latest_versions_by_environment_api_key(
+        self, environment_api_key: str
+    ) -> RawQuerySet:
+        """
+        Get the latest EnvironmentFeatureVersion objects for a given environment.
+        """
+        return self._get_latest_versions(environment_api_key=environment_api_key)
 
     def get_latest_versions_as_queryset(
-        self, environment: "Environment"
+        self, environment_id: int
     ) -> QuerySet["EnvironmentFeatureVersion"]:
         """
         Get the latest EnvironmentFeatureVersion objects for a given environment
@@ -33,5 +38,20 @@ def get_latest_versions_as_queryset(
         operations on the ORM object.
         """
         return self.filter(
-            uuid__in=[efv.uuid for efv in self.get_latest_versions(environment)]
+            uuid__in=[
+                efv.uuid
+                for efv in self._get_latest_versions(environment_id=environment_id)
+            ]
+        )
+
+    def _get_latest_versions(
+        self, environment_id: int = None, environment_api_key: str = None
+    ) -> RawQuerySet:
+        assert (environment_id or environment_api_key) and not (
+            environment_id and environment_api_key
+        ), "Must provide exactly one of environment_id or environment_api_key"
+
+        return self.raw(
+            get_latest_versions_sql,
+            params={"environment_id": environment_id, "api_key": environment_api_key},
         )

api/features/versioning/sql/get_latest_versions.sql

@@ -21,5 +21,8 @@ join (
 	efv1."feature_id" = latest_release_dates."feature_id"
 	and efv1."environment_id" = latest_release_dates."environment_id"
 	and efv1."live_from" = latest_release_dates."latest_release"
+inner join
+	environments_environment e on e.id = efv1.environment_id
 where
-	efv1.environment_id = %(environment_id)s;
+	(%(environment_id)s is not null and efv1.environment_id = %(environment_id)s)
+	or (%(api_key)s is not null and e.api_key = %(api_key)s);

api/tests/unit/environments/test_unit_environments_views.py

@@ -20,6 +20,7 @@
 from environments.permissions.constants import VIEW_ENVIRONMENT
 from environments.permissions.models import UserEnvironmentPermission
 from features.models import Feature, FeatureState
+from features.versioning.models import EnvironmentFeatureVersion
 from metadata.models import Metadata, MetadataModelField
 from organisations.models import Organisation
 from projects.models import Project
@@ -935,3 +936,27 @@ def test_cannot_enable_v2_versioning_for_environment_already_enabled(
     assert response.json() == {"detail": "Environment already using v2 versioning."}
 
     mock_enable_v2_versioning.delay.assert_not_called()
+
+
+def test_total_segment_overrides_correctly_ignores_old_versions(
+    feature: Feature,
+    segment_featurestate: FeatureState,
+    environment_v2_versioning: Environment,
+    admin_client_new: APIClient,
+    staff_user: FFAdminUser,
+) -> None:
+    # Given
+    url = reverse(
+        "api-v1:environments:environment-detail",
+        args=[environment_v2_versioning.api_key],
+    )
+
+    EnvironmentFeatureVersion.objects.create(
+        feature=feature, environment=environment_v2_versioning
+    ).publish(staff_user)
+
+    # When
+    response = admin_client_new.get(url)
+
+    # Then
+    assert response.json()["total_segment_overrides"] == 1

api/util/mappers/engine.py

@@ -205,8 +205,8 @@ def map_environment_to_engine(
         latest_environment_feature_version_uuids=(
             {
                 efv.uuid
-                for efv in EnvironmentFeatureVersion.objects.get_latest_versions(
-                    environment
+                for efv in EnvironmentFeatureVersion.objects.get_latest_versions_by_environment_id(
+                    environment.id
                 )
             }
             if environment.use_v2_feature_versioning