feat(limits): Add limits to features, segments and segment overrides (#…

…2480)

Co-authored-by: Matthew Elwell <[email protected]>

api/features/feature_segments/serializers.py

@@ -11,6 +11,17 @@ class Meta:
 
     def validate(self, data):
         data = super().validate(data)
+        environment = data["environment"]
+        if (
+            environment.feature_segments.count()
+            >= environment.project.max_segment_overrides_allowed
+        ):
+            raise serializers.ValidationError(
+                {
+                    "environment": "The environment has reached the maximum allowed segments overrides limit."
+                }
+            )
+
         segment = data["segment"]
         if segment.feature is not None and segment.feature != data["feature"]:
             raise serializers.ValidationError(

api/features/feature_segments/tests/test_views.py

@@ -271,3 +271,38 @@ def test_creating_segment_override_for_feature_based_segment_returns_201_for_cor
     response = client.post(url, data=json.dumps(data), content_type="application/json")
     # Then
     assert response.status_code == status.HTTP_201_CREATED
+
+
+@pytest.mark.parametrize(
+    "client", [lazy_fixture("master_api_key_client"), lazy_fixture("admin_client")]
+)
+def test_creating_segment_override_reaching_max_limit(
+    client, segment, environment, project, feature, feature_based_segment
+):
+    # Given
+    project.max_segment_overrides_allowed = 1
+    project.save()
+
+    data = {
+        "feature": feature.id,
+        "segment": segment.id,
+        "environment": environment.id,
+    }
+    url = reverse("api-v1:features:feature-segment-list")
+    # let's create the first segment override
+    response = client.post(url, data=json.dumps(data), content_type="application/json")
+    assert response.status_code == status.HTTP_201_CREATED
+
+    # Then - Try to create another override
+    data = {
+        "feature": feature.id,
+        "segment": feature_based_segment.id,
+        "environment": environment.id,
+    }
+    response = client.post(url, data=json.dumps(data), content_type="application/json")
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert (
+        response.json()["environment"][0]
+        == "The environment has reached the maximum allowed segments overrides limit."
+    )
+    assert environment.feature_segments.count() == 1

api/features/serializers.py

@@ -6,9 +6,11 @@
 from rest_framework.exceptions import PermissionDenied
 
 from environments.identities.models import Identity
+from environments.models import Environment
 from environments.sdk.serializers_mixins import (
     HideSensitiveFieldsSerializerMixin,
 )
+from projects.models import Project
 from users.serializers import UserIdsSerializer, UserListSerializer
 from util.drf_writable_nested.serializers import (
     DeleteBeforeUpdateWritableNestedModelSerializer,
@@ -119,7 +121,10 @@ def to_internal_value(self, data):
             data["initial_value"] = str(data["initial_value"])
         return super(ListCreateFeatureSerializer, self).to_internal_value(data)
 
-    def create(self, validated_data):
+    def create(self, validated_data: dict) -> Feature:
+        project = self.context["project"]
+        self.validate_project_features_limit(project)
+
         # Add the default(User creating the feature) owner of the feature
         # NOTE: pop the user before passing the data to create
         user = validated_data.pop("user", None)
@@ -128,6 +133,14 @@ def create(self, validated_data):
             instance.owners.add(user)
         return instance
 
+    def validate_project_features_limit(self, project: Project) -> None:
+        if project.features.count() >= project.max_features_allowed:
+            raise serializers.ValidationError(
+                {
+                    "project": "The Project has reached the maximum allowed features limit."
+                }
+            )
+
     def validate_multivariate_options(self, multivariate_options):
         if multivariate_options:
             user = self.context["request"].user
@@ -427,3 +440,21 @@ def _get_save_kwargs(self, field_name):
             kwargs["feature"] = self.context.get("feature")
             kwargs["environment"] = self.context.get("environment")
         return kwargs
+
+    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
+        ):
+            raise serializers.ValidationError(
+                {
+                    "environment": "The environment has reached the maximum allowed segments overrides limit."
+                }
+            )

api/projects/admin.py

@@ -63,4 +63,7 @@ class ProjectAdmin(admin.ModelAdmin):
         "hide_disabled_flags",
         "enable_dynamo_db",
         "enable_realtime_updates",
+        "max_segments_allowed",
+        "max_features_allowed",
+        "max_segment_overrides_allowed",
     )

api/projects/exceptions.py

@@ -14,3 +14,8 @@ class ProjectMigrationError(APIException):
 class TooManyIdentitiesError(APIException):
     status_code = 400
     default_detail = "Too many identities; Please contact support"
+
+
+class ProjectTooLargeError(APIException):
+    status_code = 400
+    default_detail = "Project is too large; Please contact support"

api/projects/migrations/0019_add_limits.py

@@ -0,0 +1,34 @@
+# Generated by Django 3.2.20 on 2023-07-21 08:26
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("projects", "0018_add_unique_permission_constraint"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="project",
+            name="max_features_allowed",
+            field=models.IntegerField(
+                default=400, help_text="Max features allowed for this project"
+            ),
+        ),
+        migrations.AddField(
+            model_name="project",
+            name="max_segments_allowed",
+            field=models.IntegerField(
+                default=100, help_text="Max segments allowed for this project"
+            ),
+        ),
+        migrations.AddField(
+            model_name="project",
+            name="max_segment_overrides_allowed",
+            field=models.IntegerField(
+                default=100,
+                help_text="Max segments overrides allowed for any (one) environment within this project",
+            ),
+        ),
+    ]

api/projects/models.py

@@ -7,6 +7,7 @@
 from django.conf import settings
 from django.core.cache import caches
 from django.db import models
+from django.db.models import Count
 from django.utils import timezone
 from django_lifecycle import (
     AFTER_SAVE,
@@ -60,6 +61,16 @@ class Project(LifecycleModelMixin, SoftDeleteExportableModel):
         max_length=255,
         help_text="Used for validating feature names",
     )
+    max_segments_allowed = models.IntegerField(
+        default=100, help_text="Max segments allowed for this project"
+    )
+    max_features_allowed = models.IntegerField(
+        default=400, help_text="Max features allowed for this project"
+    )
+    max_segment_overrides_allowed = models.IntegerField(
+        default=100,
+        help_text="Max segments overrides allowed for any (one) environment within this project",
+    )
 
     objects = ProjectManager()
 
@@ -69,6 +80,18 @@ class Meta:
     def __str__(self):
         return "Project %s" % self.name
 
+    @property
+    def is_too_large(self) -> bool:
+        return (
+            self.features.count() > self.max_features_allowed
+            or self.segments.count() > self.max_segments_allowed
+            or self.environments.annotate(
+                segment_override_count=Count("feature_segments")
+            )
+            .filter(segment_override_count__gt=self.max_segment_overrides_allowed)
+            .exists()
+        )
+
     def get_segments_from_cache(self):
         segments = project_segments_cache.get(self.id)
 

api/projects/tests/test_views.py

@@ -11,6 +11,7 @@
 
 from environments.dynamodb.types import ProjectIdentityMigrationStatus
 from environments.identities.models import Identity
+from features.models import FeatureSegment
 from organisations.models import Organisation, OrganisationRole
 from organisations.permissions.models import (
     OrganisationPermissionModel,
@@ -534,6 +535,88 @@ def test_project_migrate_to_edge_returns_400_if_project_have_too_many_identities
     mocked_identity_migrator.assert_not_called()
 
 
+def test_project_migrate_to_edge_returns_400_if_project_have_too_many_features(
+    admin_client, project, mocker, environment, feature, multivariate_feature, settings
+):
+    # Given
+    settings.PROJECT_METADATA_TABLE_NAME_DYNAMO = "some_table"
+    project.max_features_allowed = 1
+    project.save()
+
+    mocked_identity_migrator = mocker.patch("projects.views.IdentityMigrator")
+
+    url = reverse("api-v1:projects:project-migrate-to-edge", args=[project.id])
+
+    # When
+    response = admin_client.post(url)
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.json()["detail"] == "Project is too large; Please contact support"
+    mocked_identity_migrator.assert_not_called()
+
+
+def test_project_migrate_to_edge_returns_400_if_project_have_too_many_segments(
+    admin_client,
+    project,
+    mocker,
+    environment,
+    feature,
+    settings,
+    feature_based_segment,
+    segment,
+):
+    # Given
+    settings.PROJECT_METADATA_TABLE_NAME_DYNAMO = "some_table"
+    project.max_segments_allowed = 1
+    project.save()
+
+    mocked_identity_migrator = mocker.patch("projects.views.IdentityMigrator")
+
+    url = reverse("api-v1:projects:project-migrate-to-edge", args=[project.id])
+
+    # When
+    response = admin_client.post(url)
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.json()["detail"] == "Project is too large; Please contact support"
+    mocked_identity_migrator.assert_not_called()
+
+
+def test_project_migrate_to_edge_returns_400_if_project_have_too_many_segment_overrides(
+    admin_client,
+    project,
+    mocker,
+    environment,
+    feature,
+    settings,
+    feature_segment,
+    multivariate_feature,
+    segment,
+):
+    # Given
+    settings.PROJECT_METADATA_TABLE_NAME_DYNAMO = "some_table"
+    project.max_segment_overrides_allowed = 1
+    project.save()
+
+    # let's create another feature segment
+    FeatureSegment.objects.create(
+        feature=multivariate_feature, segment=segment, environment=environment
+    )
+    mocked_identity_migrator = mocker.patch("projects.views.IdentityMigrator")
+
+    url = reverse("api-v1:projects:project-migrate-to-edge", args=[project.id])
+
+    # When
+    response = admin_client.post(url)
+
+    # Then
+    assert response.status_code == status.HTTP_400_BAD_REQUEST
+    assert response.json()["detail"] == "Project is too large; Please contact support"
+    mocked_identity_migrator.assert_not_called()
+
+
 def test_list_project_with_uuid_filter_returns_correct_project(
     admin_client, project, mocker, settings, organisation
 ):

api/projects/views.py

@@ -24,6 +24,7 @@
 from projects.exceptions import (
     DynamoNotEnabledError,
     ProjectMigrationError,
+    ProjectTooLargeError,
     TooManyIdentitiesError,
 )
 from projects.models import (
@@ -163,6 +164,9 @@ def migrate_to_edge(self, request: Request, pk: int = None):
             raise DynamoNotEnabledError()
 
         project = self.get_object()
+        if project.is_too_large:
+            raise ProjectTooLargeError()
+
         identity_count = Identity.objects.filter(environment__project=project).count()
 
         if identity_count > settings.MAX_SELF_MIGRATABLE_IDENTITIES:

api/segments/serializers.py

@@ -5,6 +5,7 @@
 from rest_framework.serializers import ListSerializer
 from rest_framework_recursive.fields import RecursiveField
 
+from projects.models import Project
 from segments.models import PERCENTAGE_SPLIT, Condition, Segment, SegmentRule
 
 
@@ -52,19 +53,26 @@ def validate(self, attrs):
         return attrs
 
     def create(self, validated_data):
-        """
-        Override create method to create segment with nested rules and conditions
+        project = validated_data["project"]
+        self.validate_project_segment_limit(project)
 
-        :param validated_data: validated json data
-        :return: created Segment object
-        """
         rules_data = validated_data.pop("rules", [])
+
+        # create segment with nested rules and conditions
         segment = Segment.objects.create(**validated_data)
         self._update_or_create_segment_rules(
             rules_data, segment=segment, is_create=True
         )
         return segment
 
+    def validate_project_segment_limit(self, project: Project) -> None:
+        if project.segments.count() >= project.max_segments_allowed:
+            raise ValidationError(
+                {
+                    "project": "The project has reached the maximum allowed segments limit."
+                }
+            )
+
     def update(self, instance, validated_data):
         # use the initial data since we need the ids included to determine which to update & which to create
         rules_data = self.initial_data.pop("rules", [])