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

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 is_feature_enabled
 from projects.code_references.services import (
     annotate_feature_queryset_with_code_references_summary,
 )
@@ -217,9 +220,13 @@ 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
+        if is_feature_enabled("code_references_ui_stats", project.organisation):
+            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/integrations/flagsmith/client.py

@@ -10,15 +10,22 @@
 ```
 """
 
+import logging
 import typing
 
 from django.conf import settings
 from flagsmith import Flagsmith
+from flagsmith.exceptions import FlagsmithFeatureDoesNotExistError
 from flagsmith.offline_handlers import LocalFileHandler
 
 from integrations.flagsmith.exceptions import FlagsmithIntegrationError
 from integrations.flagsmith.flagsmith_service import ENVIRONMENT_JSON_PATH
 
+if typing.TYPE_CHECKING:
+    from organisations.models import Organisation
+
+logger = logging.getLogger(__name__)
+
 _flagsmith_clients: dict[str, Flagsmith] = {}
 
 
@@ -36,6 +43,23 @@ def get_client(name: str = "default", local_eval: bool = False) -> Flagsmith:
     return _flagsmith_client
 
 
+def is_feature_enabled(
+    feature_name: str,
+    organisation: "Organisation",
+) -> bool:
+    """Check if a Flagsmith-on-Flagsmith feature flag is enabled for an organisation."""
+    client = get_client("local", local_eval=True)
+    flags = client.get_identity_flags(
+        organisation.flagsmith_identifier,
+        traits=organisation.flagsmith_on_flagsmith_api_traits,
+    )
+    try:
+        return flags.is_feature_enabled(feature_name)
+    except FlagsmithFeatureDoesNotExistError:
+        logger.warning("FoF feature %r not found, defaulting to disabled", feature_name)
+        return False
+
+
 def _get_client_kwargs() -> dict[str, typing.Any]:
     _default_kwargs = {"offline_handler": LocalFileHandler(ENVIRONMENT_JSON_PATH)}
 

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

@@ -3586,10 +3586,11 @@ def test_FeatureViewSet_list__includes_code_references_counts(
     project: Project,
     feature: Feature,
     with_project_permissions: WithProjectPermissionsCallable,
-    environment: Environment,
+    mocker: MockerFixture,
 ) -> None:
     # Given
     with_project_permissions([VIEW_PROJECT])  # type: ignore[call-arg]
+    mocker.patch("features.views.is_feature_enabled", return_value=True)
     with freeze_time("2099-01-01T10:00:00-0300"):
         FeatureFlagCodeReferencesScan.objects.create(
             project=project,
@@ -3667,26 +3668,62 @@ def test_FeatureViewSet_list__includes_code_references_counts(
     ]
 
 
+@pytest.mark.usefixtures("feature")
 def test_FeatureViewSet_list__no_scans__returns_empty_code_references_counts(
+    staff_client: APIClient,
+    project: Project,
+    environment: Environment,
+    with_project_permissions: WithProjectPermissionsCallable,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    with_project_permissions([VIEW_PROJECT])  # type: ignore[call-arg]
+    mocker.patch("features.views.is_feature_enabled", return_value=True)
+
+    # 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_FeatureViewSet_list__code_references_ui_stats_disabled__returns_empty_counts(
     staff_client: APIClient,
     project: Project,
     feature: Feature,
     environment: Environment,
     with_project_permissions: WithProjectPermissionsCallable,
 ) -> None:
-    # Given - project has no code reference scans
+    # Given
     with_project_permissions([VIEW_PROJECT])  # type: ignore[call-arg]
+    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"] == []
 
 

api/tests/unit/integrations/flagsmith/test_unit_flagsmith_client.py

@@ -5,9 +5,10 @@
 from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
 
-from integrations.flagsmith.client import get_client
+from integrations.flagsmith.client import get_client, is_feature_enabled
 from integrations.flagsmith.exceptions import FlagsmithIntegrationError
 from integrations.flagsmith.flagsmith_service import ENVIRONMENT_JSON_PATH
+from organisations.models import Organisation
 
 
 @pytest.fixture(autouse=True)
@@ -100,3 +101,39 @@ def test_get_client_raises_value_error_if_missing_args(  # type: ignore[no-untyp
     # When
     with pytest.raises(FlagsmithIntegrationError):
         get_client()
+
+
+@pytest.mark.parametrize("expected", [True, False])
+def test_is_feature_enabled__mocked_client__returns_expected(
+    organisation: Organisation,
+    mocker: MockerFixture,
+    expected: bool,
+) -> None:
+    # Given
+    get_client = mocker.patch("integrations.flagsmith.client.get_client")
+    get_client().get_identity_flags().is_feature_enabled.return_value = expected
+
+    # When
+    result = is_feature_enabled("some_feature", organisation)
+
+    # Then
+    assert result is expected
+
+
+def test_is_feature_enabled__unknown_feature__returns_false(
+    organisation: Organisation,
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    from flagsmith.exceptions import FlagsmithFeatureDoesNotExistError
+
+    get_client = mocker.patch("integrations.flagsmith.client.get_client")
+    get_client().get_identity_flags().is_feature_enabled.side_effect = (
+        FlagsmithFeatureDoesNotExistError("Feature does not exist: unknown_feature")
+    )
+
+    # When
+    result = is_feature_enabled("unknown_feature", organisation)
+
+    # Then
+    assert result is False

frontend/web/components/feature-summary/FeatureTags.tsx

@@ -1,15 +1,15 @@
 import React, { FC, useMemo } from 'react'
 import SegmentOverridesIcon from 'components/SegmentOverridesIcon'
 import Constants from 'common/constants'
-import { ProjectFlag, VCSProvider } from 'common/types/responses'
+import { ProjectFlag } from 'common/types/responses'
 import IdentityOverridesIcon from 'components/IdentityOverridesIcon'
 import TagValues from 'components/tags/TagValues'
 import UnhealthyFlagWarning from './UnhealthyFlagWarning'
 import StaleFlagWarning from './StaleFlagWarning'
 import Tag from 'components/tags/Tag'
 import Utils from 'common/utils/utils'
 import { useGetHealthEventsQuery } from 'common/services/useHealthEvents'
-import VCSProviderTag from 'components/tags/VCSProviderTag'
+import CodeReferencesTag from 'components/tags/CodeReferencesTag'
 
 type FeatureTagsType = {
   editFeature: (tab?: string) => void
@@ -38,14 +38,6 @@ const FeatureTags: FC<FeatureTagsType> = ({ editFeature, projectFlag }) => {
   }
   const isFeatureHealthEnabled = Utils.getFlagsmithHasFeature('feature_health')
 
-  const hasScannedCodeReferences =
-    projectFlag?.code_references_counts?.length > 0
-  const codeReferencesCounts =
-    projectFlag?.code_references_counts?.reduce(
-      (acc, curr) => acc + curr.count,
-      0,
-    ) || 0
-
   return (
     <>
       <SegmentOverridesIcon
@@ -63,28 +55,13 @@ const FeatureTags: FC<FeatureTagsType> = ({ editFeature, projectFlag }) => {
         count={projectFlag.num_identity_overrides}
         showPlusIndicator={showPlusIndicator}
       />
-      {hasScannedCodeReferences && (
-        <Tooltip
-          title={
-            <div
-              onClick={(e) => {
-                e.stopPropagation()
-                editFeature(Constants.featurePanelTabs.USAGE)
-              }}
-              style={{ cursor: 'pointer' }}
-            >
-              <VCSProviderTag
-                count={codeReferencesCounts}
-                isWarning={codeReferencesCounts === 0}
-                vcsProvider={VCSProvider.GITHUB}
-              />
-            </div>
-          }
-          place='top'
-        >
-          {`Scanned ${codeReferencesCounts?.toString()} times in ${projectFlag?.code_references_counts?.length?.toString()} repositories`}
-        </Tooltip>
-      )}
+      <CodeReferencesTag
+        projectFlag={projectFlag}
+        onClick={(e) => {
+          e.stopPropagation()
+          editFeature(Constants.featurePanelTabs.USAGE)
+        }}
+      />
       {projectFlag.is_server_key_only && (
         <Tooltip
           title={

frontend/web/components/feature-summary/ProjectFeatureRow.tsx

@@ -1,10 +1,10 @@
 import React, { FC } from 'react'
 import classNames from 'classnames'
-import { ProjectFlag, VCSProvider } from 'common/types/responses'
+import { ProjectFlag } from 'common/types/responses'
 import FeatureName from './FeatureName'
 import FeatureDescription from './FeatureDescription'
 import TagValues from 'components/tags/TagValues'
-import VCSProviderTag from 'components/tags/VCSProviderTag'
+import CodeReferencesTag from 'components/tags/CodeReferencesTag'
 
 export interface ProjectFeatureRowProps {
   projectFlag: ProjectFlag
@@ -25,14 +25,6 @@ const ProjectFeatureRow: FC<ProjectFeatureRowProps> = ({
 }) => {
   const { description } = projectFlag
 
-  const hasScannedCodeReferences =
-    projectFlag?.code_references_counts?.length > 0
-  const codeReferencesCounts =
-    projectFlag?.code_references_counts?.reduce(
-      (acc, curr) => acc + curr.count,
-      0,
-    ) || 0
-
   return (
     <>
       {/* Desktop */}
@@ -61,20 +53,7 @@ const ProjectFeatureRow: FC<ProjectFeatureRowProps> = ({
           <div className='mx-0 flex-1 flex-column'>
             <div className='d-flex align-items-center'>
               <FeatureName name={projectFlag.name} />
-              {hasScannedCodeReferences && (
-                <Tooltip
-                  title={
-                    <VCSProviderTag
-                      count={codeReferencesCounts}
-                      isWarning={codeReferencesCounts === 0}
-                      vcsProvider={VCSProvider.GITHUB}
-                    />
-                  }
-                  place='top'
-                >
-                  {`Scanned ${codeReferencesCounts} times in ${projectFlag?.code_references_counts?.length} repositories`}
-                </Tooltip>
-              )}
+              <CodeReferencesTag projectFlag={projectFlag} />
               <TagValues
                 projectId={`${projectFlag.project}`}
                 value={projectFlag.tags}
@@ -110,20 +89,7 @@ const ProjectFeatureRow: FC<ProjectFeatureRowProps> = ({
           )}
           <div className='flex-1 align-items-center flex-wrap'>
             <FeatureName name={projectFlag.name} />
-            {hasScannedCodeReferences && (
-              <Tooltip
-                title={
-                  <VCSProviderTag
-                    count={codeReferencesCounts}
-                    isWarning={codeReferencesCounts === 0}
-                    vcsProvider={VCSProvider.GITHUB}
-                  />
-                }
-                place='top'
-              >
-                {`Scanned ${codeReferencesCounts} times in ${projectFlag?.code_references_counts?.length} repositories`}
-              </Tooltip>
-            )}
+            <CodeReferencesTag projectFlag={projectFlag} />
             <TagValues
               projectId={`${projectFlag.project}`}
               value={projectFlag.tags}