fix: ignore old versions when validating segment override limit (#4618)

api/features/feature_segments/limits.py

@@ -0,0 +1,49 @@
+from django.db.models import Q
+
+from environments.models import Environment
+from features.versioning.models import EnvironmentFeatureVersion
+
+SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE = (
+    "The environment has reached the maximum allowed segments overrides limit."
+)
+
+
+def exceeds_segment_override_limit(
+    environment: Environment,
+    segment_ids_to_create_overrides: list[int] | None = None,
+    segment_ids_to_delete_overrides: list[int] | None = None,
+    exclusive: bool = False,
+) -> bool:
+    q = Q()
+
+    segment_ids_to_create_overrides = segment_ids_to_create_overrides or []
+    segment_ids_to_delete_overrides = segment_ids_to_delete_overrides or []
+
+    def _check(left: int, right: int) -> bool:
+        if exclusive:
+            return left > right
+        return left >= right
+
+    if environment.use_v2_feature_versioning:
+        latest_versions = (
+            EnvironmentFeatureVersion.objects.get_latest_versions_by_environment_id(
+                environment_id=environment.id
+            )
+        )
+        q = q & Q(environment_feature_version__in=latest_versions)
+
+    existing_overridden_segment_ids = set(
+        environment.feature_segments.filter(q).values_list("segment_id", flat=True)
+    )
+    segment_override_count = len(existing_overridden_segment_ids)
+
+    extra = len(segment_ids_to_create_overrides) - len(
+        set(segment_ids_to_delete_overrides).intersection(
+            existing_overridden_segment_ids
+        )
+    )
+
+    return _check(
+        segment_override_count + extra,
+        environment.project.max_segment_overrides_allowed,
+    )

api/features/feature_segments/serializers.py

@@ -6,6 +6,10 @@
 from rest_framework.exceptions import PermissionDenied
 
 from environments.permissions.constants import MANAGE_SEGMENT_OVERRIDES
+from features.feature_segments.limits import (
+    SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE,
+    exceeds_segment_override_limit,
+)
 from features.models import FeatureSegment
 
 
@@ -17,15 +21,9 @@ class Meta:
 
     def validate(self, data):
         data = super().validate(data)
-        environment = data["environment"]
-        if (
-            environment.feature_segments.count()
-            >= environment.project.max_segment_overrides_allowed
-        ):
+        if exceeds_segment_override_limit(data["environment"]):
             raise serializers.ValidationError(
-                {
-                    "environment": "The environment has reached the maximum allowed segments overrides limit."
-                }
+                {"environment": SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE}
             )
 
         segment = data["segment"]

api/features/serializers.py

@@ -15,7 +15,6 @@
 from rest_framework.exceptions import PermissionDenied
 
 from environments.identities.models import Identity
-from environments.models import Environment
 from environments.sdk.serializers_mixins import (
     HideSensitiveFieldsSerializerMixin,
 )
@@ -33,6 +32,10 @@
 )
 
 from .constants import INTERSECTION, UNION
+from .feature_segments.limits import (
+    SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE,
+    exceeds_segment_override_limit,
+)
 from .feature_segments.serializers import (
     CustomCreateSegmentOverrideFeatureSegmentSerializer,
 )
@@ -599,6 +602,8 @@ class SDKFeatureStatesQuerySerializer(serializers.Serializer):
 class CustomCreateSegmentOverrideFeatureStateSerializer(
     CreateSegmentOverrideFeatureStateSerializer
 ):
+    validate_override_limit = True
+
     feature_segment = CustomCreateSegmentOverrideFeatureSegmentSerializer(
         required=False, allow_null=True
     )
@@ -615,18 +620,8 @@ def _get_save_kwargs(self, field_name):
 
     def create(self, validated_data: dict) -> FeatureState:
         environment = validated_data["environment"]
-        self.validate_environment_segment_override_limit(environment)
-        return super().create(validated_data)
-
-    def validate_environment_segment_override_limit(
-        self, environment: Environment
-    ) -> None:
-        if (
-            environment.feature_segments.count()
-            >= environment.project.max_segment_overrides_allowed
-        ):
+        if self.validate_override_limit and exceeds_segment_override_limit(environment):
             raise serializers.ValidationError(
-                {
-                    "environment": "The environment has reached the maximum allowed segments overrides limit."
-                }
+                {"environment": SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE}
             )
+        return super().create(validated_data)

api/features/versioning/serializers.py

@@ -5,6 +5,11 @@
 from rest_framework import serializers
 
 from api_keys.user import APIKeyUser
+from environments.models import Environment
+from features.feature_segments.limits import (
+    SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE,
+    exceeds_segment_override_limit,
+)
 from features.serializers import (
     CustomCreateSegmentOverrideFeatureStateSerializer,
 )
@@ -18,6 +23,8 @@
 class CustomEnvironmentFeatureVersionFeatureStateSerializer(
     CustomCreateSegmentOverrideFeatureStateSerializer
 ):
+    validate_override_limit = False
+
     class Meta(CustomCreateSegmentOverrideFeatureStateSerializer.Meta):
         read_only_fields = (
             CustomCreateSegmentOverrideFeatureStateSerializer.Meta.read_only_fields
@@ -141,10 +148,24 @@ class Meta(EnvironmentFeatureVersionSerializer.Meta):
             "publish_immediately",
         )
 
+    @property
+    def segment_ids_to_create_overrides(self) -> list[int]:
+        return [
+            fs["feature_segment"]["segment"]
+            for fs in self.initial_data.get("feature_states_to_create", [])
+            if fs["feature_segment"] is not None
+        ]
+
     @transaction.atomic
     def create(
         self, validated_data: dict[str, typing.Any]
     ) -> EnvironmentFeatureVersion:
+        # Since the environment is passed as a keyword argument to save, we have
+        # to perform this validation here instead of in the `validate` method.
+        environment = validated_data["environment"]
+        self._validate_segment_override_limit(environment)
+        self._validate_v2_versioning(environment)
+
         # Note that we use self.initial_data below for handling the feature states
         # since we want the raw data (rather than the serialized ORM objects) to pass
         # into the serializers in the separate private methods used for modifying the
@@ -284,6 +305,25 @@ def _delete_feature_states(
     def _is_segment_override(self, feature_state: dict) -> bool:
         return feature_state.get("feature_segment") is not None
 
+    def _validate_segment_override_limit(self, environment: Environment) -> None:
+        if exceeds_segment_override_limit(
+            environment=environment,
+            segment_ids_to_create_overrides=self.segment_ids_to_create_overrides,
+            segment_ids_to_delete_overrides=self.initial_data.get(
+                "segment_ids_to_delete_overrides", []
+            ),
+            exclusive=True,
+        ):
+            raise serializers.ValidationError(
+                {"environment": SEGMENT_OVERRIDE_LIMIT_EXCEEDED_MESSAGE}
+            )
+
+    def _validate_v2_versioning(self, environment: Environment) -> None:
+        if not environment.use_v2_feature_versioning:
+            raise serializers.ValidationError(
+                {"environment": "Environment must use v2 feature versioning."}
+            )
+
 
 class EnvironmentFeatureVersionPublishSerializer(serializers.Serializer):
     live_from = serializers.DateTimeField(required=False)

api/tests/unit/features/feature_segments/test_unit_feature_segments_limits.py

@@ -0,0 +1,27 @@
+from environments.models import Environment
+from features.feature_segments.limits import exceeds_segment_override_limit
+from features.models import Feature
+from projects.models import Project
+from segments.models import Segment
+
+
+def test_segment_override_limit_does_not_exclude_invalid_overrides_being_deleted(
+    feature: Feature,
+    segment: Segment,
+    another_segment: Segment,
+    environment_v2_versioning: Environment,
+    project: Project,
+) -> None:
+    # Given
+    project.max_segment_overrides_allowed = 0
+    project.save()
+
+    # When
+    result = exceeds_segment_override_limit(
+        environment=environment_v2_versioning,
+        segment_ids_to_create_overrides=[another_segment.id],
+        segment_ids_to_delete_overrides=[segment.id],
+    )
+
+    # Then
+    assert result is True