fix: Enable faster feature loading (#3550)

api/features/serializers.py

@@ -3,6 +3,7 @@
 
 import django.core.exceptions
 from drf_writable_nested import WritableNestedModelSerializer
+from drf_yasg.utils import swagger_serializer_method
 from rest_framework import serializers
 from rest_framework.exceptions import PermissionDenied
 
@@ -32,47 +33,22 @@
 )
 
 
-class FeatureOwnerInputSerializer(UserIdsSerializer):
-    def add_owners(self, feature: Feature):
-        user_ids = self.validated_data["user_ids"]
-        feature.owners.add(*user_ids)
-
-    def remove_users(self, feature: Feature):
-        user_ids = self.validated_data["user_ids"]
-        feature.owners.remove(*user_ids)
-
-
-class FeatureGroupOwnerInputSerializer(serializers.Serializer):
-    group_ids = serializers.ListField(child=serializers.IntegerField())
-
-    def add_group_owners(self, feature: Feature):
-        group_ids = self.validated_data["group_ids"]
-        feature.group_owners.add(*group_ids)
-
-    def remove_group_owners(self, feature: Feature):
-        group_ids = self.validated_data["group_ids"]
-        feature.group_owners.remove(*group_ids)
-
-
-class ProjectFeatureSerializer(serializers.ModelSerializer):
-    owners = UserListSerializer(many=True, read_only=True)
-    group_owners = UserPermissionGroupSummarySerializer(many=True, read_only=True)
+class FeatureStateSerializerSmall(serializers.ModelSerializer):
+    feature_state_value = serializers.SerializerMethodField()
 
     class Meta:
-        model = Feature
+        model = FeatureState
         fields = (
             "id",
-            "name",
-            "created_date",
-            "description",
-            "initial_value",
-            "default_enabled",
-            "type",
-            "owners",
-            "group_owners",
-            "is_server_key_only",
+            "feature_state_value",
+            "environment",
+            "identity",
+            "feature_segment",
+            "enabled",
         )
-        writeonly_fields = ("initial_value", "default_enabled")
+
+    def get_feature_state_value(self, obj):
+        return obj.get_feature_state_value(identity=self.context.get("identity"))
 
 
 class FeatureQuerySerializer(serializers.Serializer):
@@ -147,6 +123,9 @@ class CreateFeatureSerializer(DeleteBeforeUpdateWritableNestedModelSerializer):
     )
     owners = UserListSerializer(many=True, read_only=True)
     group_owners = UserPermissionGroupSummarySerializer(many=True, read_only=True)
+
+    environment_feature_state = serializers.SerializerMethodField()
+
     num_segment_overrides = serializers.SerializerMethodField(
         help_text="Number of segment overrides that exist for the given feature "
         "in the environment provided by the `environment` query parameter."
@@ -186,6 +165,7 @@ class Meta:
             "group_owners",
             "uuid",
             "project",
+            "environment_feature_state",
             "num_segment_overrides",
             "num_identity_overrides",
             "is_server_key_only",
@@ -282,6 +262,17 @@ def validate(self, attrs):
 
         return attrs
 
+    @swagger_serializer_method(
+        serializer_or_field=FeatureStateSerializerSmall(allow_null=True)
+    )
+    def get_environment_feature_state(
+        self, instance: Feature
+    ) -> dict[str, typing.Any] | None:
+        if (feature_states := self.context.get("feature_states")) and (
+            feature_state := feature_states.get(instance.id)
+        ):
+            return FeatureStateSerializerSmall(instance=feature_state).data
+
     def get_num_segment_overrides(self, instance) -> int:
         try:
             return self.context["overrides_data"][instance.id].num_segment_overrides
@@ -366,6 +357,49 @@ def get_feature_state_value(self, obj):
         return obj.get_feature_state_value(identity=self.context.get("identity"))
 
 
+class FeatureOwnerInputSerializer(UserIdsSerializer):
+    def add_owners(self, feature: Feature):
+        user_ids = self.validated_data["user_ids"]
+        feature.owners.add(*user_ids)
+
+    def remove_users(self, feature: Feature):
+        user_ids = self.validated_data["user_ids"]
+        feature.owners.remove(*user_ids)
+
+
+class FeatureGroupOwnerInputSerializer(serializers.Serializer):
+    group_ids = serializers.ListField(child=serializers.IntegerField())
+
+    def add_group_owners(self, feature: Feature):
+        group_ids = self.validated_data["group_ids"]
+        feature.group_owners.add(*group_ids)
+
+    def remove_group_owners(self, feature: Feature):
+        group_ids = self.validated_data["group_ids"]
+        feature.group_owners.remove(*group_ids)
+
+
+class ProjectFeatureSerializer(serializers.ModelSerializer):
+    owners = UserListSerializer(many=True, read_only=True)
+    group_owners = UserPermissionGroupSummarySerializer(many=True, read_only=True)
+
+    class Meta:
+        model = Feature
+        fields = (
+            "id",
+            "name",
+            "created_date",
+            "description",
+            "initial_value",
+            "default_enabled",
+            "type",
+            "owners",
+            "group_owners",
+            "is_server_key_only",
+        )
+        writeonly_fields = ("initial_value", "default_enabled")
+
+
 class SDKFeatureStateSerializer(
     HideSensitiveFieldsSerializerMixin, FeatureStateSerializerFull
 ):

api/features/views.py

@@ -158,8 +158,31 @@ def get_queryset(self):
         )
         queryset = queryset.order_by(sort)
 
+        if environment_id:
+            page = self.paginate_queryset(queryset)
+
+            self.environment = Environment.objects.get(id=environment_id)
+            q = Q(
+                feature_id__in=[feature.id for feature in page],
+                identity__isnull=True,
+                feature_segment__isnull=True,
+            )
+            feature_states = FeatureState.objects.get_live_feature_states(
+                self.environment,
+                additional_filters=q,
+            ).select_related("feature_state_value", "feature")
+
+            self._feature_states = {fs.feature_id: fs for fs in feature_states}
+
         return queryset
 
+    def paginate_queryset(self, queryset: QuerySet[Feature]) -> list[Feature]:
+        if getattr(self, "_page", None):
+            return self._page
+
+        self._page = super().paginate_queryset(queryset)
+        return self._page
+
     def perform_create(self, serializer):
         serializer.save(
             project_id=int(self.kwargs.get("project_pk")), user=self.request.user
@@ -178,8 +201,11 @@ def perform_destroy(self, instance):
     def get_serializer_context(self):
         context = super().get_serializer_context()
 
+        feature_states = getattr(self, "_feature_states", {})
         project = get_object_or_404(Project.objects.all(), pk=self.kwargs["project_pk"])
-        context.update(project=project, user=self.request.user)
+        context.update(
+            project=project, user=self.request.user, feature_states=feature_states
+        )
 
         if self.action == "list" and "environment" in self.request.query_params:
             environment = get_object_or_404(

api/tests/unit/features/test_unit_features_views.py

@@ -2402,7 +2402,7 @@ def test_list_features_n_plus_1(
         v1_feature_state.clone(env=environment, version=i, live_from=timezone.now())
 
     # When
-    with django_assert_num_queries(14):
+    with django_assert_num_queries(16):
         response = staff_client.get(url)
 
     # Then
@@ -2511,6 +2511,125 @@ def test_list_features_with_intersection_tag(
     assert response.data["results"][0]["tags"] == [tag1.id, tag2.id]
 
 
+def test_list_features_with_feature_state(
+    staff_client: APIClient,
+    project: Project,
+    feature: Feature,
+    with_project_permissions: WithProjectPermissionsCallable,
+    django_assert_num_queries: DjangoAssertNumQueries,
+    environment: Environment,
+    identity: Identity,
+    feature_segment: FeatureSegment,
+) -> None:
+    # Given
+    with_project_permissions([VIEW_PROJECT])
+
+    feature2 = Feature.objects.create(
+        name="another_feature", project=project, initial_value="initial_value"
+    )
+    feature3 = Feature.objects.create(
+        name="fancy_feature", project=project, initial_value="gone"
+    )
+
+    Environment.objects.create(
+        name="Out of test scope environment",
+        project=project,
+    )
+
+    feature_state1 = feature.feature_states.filter(environment=environment).first()
+    feature_state1.enabled = True
+    feature_state1.version = 1
+    feature_state1.save()
+
+    feature_state_value1 = feature_state1.feature_state_value
+    feature_state_value1.string_value = None
+    feature_state_value1.integer_value = 1945
+    feature_state_value1.type = INTEGER
+    feature_state_value1.save()
+
+    # This should be ignored due to versioning.
+    feature_state_versioned = FeatureState.objects.create(
+        feature=feature,
+        environment=environment,
+        enabled=True,
+        version=100,
+    )
+    feature_state_value_versioned = feature_state_versioned.feature_state_value
+    feature_state_value_versioned.string_value = None
+    feature_state_value_versioned.integer_value = 2005
+    feature_state_value_versioned.type = INTEGER
+    feature_state_value_versioned.save()
+
+    feature_state2 = feature2.feature_states.filter(environment=environment).first()
+    feature_state2.enabled = True
+    feature_state2.save()
+
+    feature_state_value2 = feature_state2.feature_state_value
+    feature_state_value2.string_value = None
+    feature_state_value2.boolean_value = True
+    feature_state_value2.type = BOOLEAN
+    feature_state_value2.save()
+
+    feature_state_value3 = (
+        feature3.feature_states.filter(environment=environment)
+        .first()
+        .feature_state_value
+    )
+    feature_state_value3.string_value = "present"
+    feature_state_value3.save()
+
+    # This should be ignored due to identity being set.
+    FeatureState.objects.create(
+        feature=feature2,
+        environment=environment,
+        identity=identity,
+    )
+
+    # This should be ignored due to feature segment being set.
+    FeatureState.objects.create(
+        feature=feature2,
+        environment=environment,
+        feature_segment=feature_segment,
+    )
+
+    # Multivariate should be ignored.
+    MultivariateFeatureOption.objects.create(
+        feature=feature2,
+        default_percentage_allocation=30,
+        type=STRING,
+        string_value="mv_feature_option1",
+    )
+    MultivariateFeatureOption.objects.create(
+        feature=feature2,
+        default_percentage_allocation=70,
+        type=STRING,
+        string_value="mv_feature_option2",
+    )
+
+    base_url = reverse("api-v1:projects:project-features-list", args=[project.id])
+    url = f"{base_url}?environment={environment.id}"
+
+    # When
+    with django_assert_num_queries(16):
+        response = staff_client.get(url)
+
+    # Then
+    assert response.status_code == status.HTTP_200_OK
+
+    assert len(response.data["results"]) == 3
+    results = response.data["results"]
+
+    assert results[0]["environment_feature_state"]["enabled"] is True
+    assert results[0]["environment_feature_state"]["feature_state_value"] == 2005
+    assert results[0]["name"] == feature.name
+    assert results[1]["environment_feature_state"]["enabled"] is True
+    assert results[1]["environment_feature_state"]["feature_state_value"] is True
+    assert results[1]["name"] == feature2.name
+    assert results[2]["environment_feature_state"]["enabled"] is False
+    assert results[2]["environment_feature_state"]["feature_state_value"] == "present"
+    assert results[2]["name"] == feature3.name
+
+
 def test_list_features_with_filter_by_value_search_string_and_int(
     staff_client: APIClient,
     project: Project,
@@ -2759,7 +2878,7 @@ def test_feature_list_last_modified_values(
         Feature.objects.create(name=f"feature_{i}", project=project)
 
     # When
-    with django_assert_num_queries(14):  # TODO: reduce this number of queries!
+    with django_assert_num_queries(16):  # TODO: reduce this number of queries!
         response = staff_client.get(url)
 
     # Then