feat(Features API): Restrict code references querying to a feature flag (#6970)

Author: emyller

api/features/views.py

@@ -6,11 +6,13 @@
 from common.core.utils import is_database_replica_setup, using_database_replica
 from common.projects.permissions import VIEW_PROJECT
 from django.conf import settings
+from django.contrib.postgres.fields import ArrayField
 from django.core.cache import caches
 from django.db.models import (
     BooleanField,
     Case,
     Exists,
+    JSONField,
     Max,
     OuterRef,
     Q,
@@ -60,6 +62,7 @@
     NestedEnvironmentPermissions,
 )
 from features.value_types import BOOLEAN, INTEGER, STRING
+from integrations.flagsmith.client import get_client
 from projects.code_references.services import (
     annotate_feature_queryset_with_code_references_summary,
 )
@@ -217,9 +220,19 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
         query_serializer.is_valid(raise_exception=True)
         query_data = query_serializer.validated_data
 
-        queryset = annotate_feature_queryset_with_code_references_summary(
-            queryset, project.id
+        # TODO: Delete this after https://github.com/flagsmith/flagsmith/issues/6832 is resolved
+        organisation = project.organisation
+        flagsmith_client = get_client("local", local_eval=True)
+        flags = flagsmith_client.get_identity_flags(
+            organisation.flagsmith_identifier,
+            traits=organisation.flagsmith_on_flagsmith_api_traits,
         )
+        if flags.is_feature_enabled("code_references_ui_stats"):
+            queryset = annotate_feature_queryset_with_code_references_summary(queryset)
+        else:
+            queryset = queryset.annotate(
+                code_references_counts=Value([], output_field=ArrayField(JSONField()))
+            )
 
         queryset = self._filter_queryset(queryset, query_serializer)
 

api/projects/code_references/services.py

@@ -2,12 +2,10 @@
 from urllib.parse import urljoin
 
 from django.contrib.postgres.expressions import ArraySubquery
-from django.contrib.postgres.fields import ArrayField
 from django.db.models import (
     BooleanField,
     F,
     Func,
-    JSONField,
     OuterRef,
     QuerySet,
     Subquery,
@@ -30,26 +28,13 @@
 
 def annotate_feature_queryset_with_code_references_summary(
     queryset: QuerySet[Feature],
-    project_id: int,
 ) -> QuerySet[Feature]:
     """Extend feature objects with a `code_references_counts`
 
     NOTE: This adds compatibility with `CodeReferenceRepositoryCountSerializer`
     while preventing N+1 queries from the serializer.
     """
     history_delta = timedelta(days=FEATURE_FLAG_CODE_REFERENCES_RETENTION_DAYS)
-    cutoff_date = timezone.now() - history_delta
-
-    # Early exit: if no scans exist for this project, skip the expensive annotation
-    has_scans = FeatureFlagCodeReferencesScan.objects.filter(
-        project_id=project_id,
-        created_at__gte=cutoff_date,
-    ).exists()
-
-    if not has_scans:
-        return queryset.annotate(
-            code_references_counts=Value([], output_field=ArrayField(JSONField()))
-        )
     last_feature_found_at = (
         FeatureFlagCodeReferencesScan.objects.annotate(
             feature_name=OuterRef("feature_name"),

api/tests/unit/features/test_unit_features_views.py

@@ -62,6 +62,7 @@
 from projects.tags.models import Tag
 from segments.models import Segment
 from tests.types import (
+    EnableFeaturesFixture,
     WithEnvironmentPermissionsCallable,
     WithProjectPermissionsCallable,
 )
@@ -3593,14 +3594,15 @@ def test_list_features__value_search_boolean__returns_matching(
 
 
 def test_list_features__with_code_references__returns_counts(
-    staff_client: APIClient,
-    project: Project,
+    enable_features: EnableFeaturesFixture,
     feature: Feature,
+    project: Project,
+    staff_client: APIClient,
     with_project_permissions: WithProjectPermissionsCallable,
-    environment: Environment,
 ) -> None:
     # Given
     with_project_permissions([VIEW_PROJECT])  # type: ignore[call-arg]
+    enable_features("code_references_ui_stats")
     with freeze_time("2099-01-01T10:00:00-0300"):
         FeatureFlagCodeReferencesScan.objects.create(
             project=project,
@@ -3678,26 +3680,64 @@ def test_list_features__with_code_references__returns_counts(
     ]
 
 
-def test_FeatureViewSet_list__no_scans__returns_empty_code_references_counts(
-    staff_client: APIClient,
+@pytest.mark.usefixtures("feature")
+def test_list_features__without_code_references__returns_empty_counts(
+    enable_features: EnableFeaturesFixture,
+    environment: Environment,
     project: Project,
-    feature: Feature,
+    staff_client: APIClient,
+    with_project_permissions: WithProjectPermissionsCallable,
+) -> None:
+    # Given
+    with_project_permissions([VIEW_PROJECT])  # type: ignore[call-arg]
+    enable_features("code_references_ui_stats")
+
+    # When
+    response = staff_client.get(
+        f"/api/v1/projects/{project.id}/features/?environment={environment.id}"
+    )
+
+    # Then
+    assert response.status_code == 200
+    results = response.json()["results"]
+    assert len(results) == 1
+    assert results[0]["code_references_counts"] == []
+
+
+# TODO: Delete this after https://github.com/flagsmith/flagsmith/issues/6832 is resolved
+def test_list_features__code_references_ui_stats_disabled__returns_empty_counts(
+    enable_features: EnableFeaturesFixture,
     environment: Environment,
+    feature: Feature,
+    project: Project,
+    staff_client: APIClient,
     with_project_permissions: WithProjectPermissionsCallable,
 ) -> None:
-    # Given - project has no code reference scans
+    # Given
     with_project_permissions([VIEW_PROJECT])  # type: ignore[call-arg]
+    enable_features()  # code_references_ui_stats not enabled
+    FeatureFlagCodeReferencesScan.objects.create(
+        project=project,
+        repository_url="https://github.flagsmith.com/backend/",
+        revision="rev-1",
+        code_references=[
+            {
+                "feature_name": feature.name,
+                "file_path": "path/to/file.py",
+                "line_number": 42,
+            },
+        ],
+    )
 
     # When
     response = staff_client.get(
         f"/api/v1/projects/{project.id}/features/?environment={environment.id}"
     )
 
-    # Then - response should include code_references_counts as empty list
+    # Then
     assert response.status_code == 200
     results = response.json()["results"]
     assert len(results) == 1
-    assert "code_references_counts" in results[0]
     assert results[0]["code_references_counts"] == []