fix(Code References): Code references are slow to query

api/features/serializers.py

@@ -443,6 +443,7 @@ def get_last_modified_in_current_environment(
 class FeatureSerializerWithMetadata(MetadataSerializerMixin, CreateFeatureSerializer):
     metadata = MetadataSerializer(required=False, many=True)
 
+    # NOTE: This field is populated by `projects.code_references.services.annotate_feature_queryset_with_code_references_summary`.
     code_references_counts = FeatureFlagCodeReferencesRepositoryCountSerializer(
         many=True,
         read_only=True,

api/features/views.py

@@ -6,13 +6,11 @@
 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,
@@ -62,7 +60,6 @@
     NestedEnvironmentPermissions,
 )
 from features.value_types import BOOLEAN, INTEGER, STRING
-from integrations.flagsmith.client import get_openfeature_client
 from projects.code_references.services import (
     annotate_feature_queryset_with_code_references_summary,
 )
@@ -219,18 +216,7 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
         query_serializer.is_valid(raise_exception=True)
         query_data = query_serializer.validated_data
 
-        # TODO: Delete this after https://github.com/flagsmith/flagsmith/issues/6832 is resolved
-        organisation = project.organisation
-        if get_openfeature_client().get_boolean_value(
-            "code_references_ui_stats",
-            default_value=False,
-            evaluation_context=organisation.openfeature_evaluation_context,
-        ):
-            queryset = annotate_feature_queryset_with_code_references_summary(queryset)
-        else:
-            queryset = queryset.annotate(
-                code_references_counts=Value([], output_field=ArrayField(JSONField()))
-            )
+        queryset = annotate_feature_queryset_with_code_references_summary(queryset)
 
         queryset = self._filter_queryset(queryset, query_serializer)
 

api/integrations/flagsmith/data/environment.json

@@ -92,19 +92,6 @@
       "featurestate_uuid": "e0d380a6-bdbc-4ad6-ae6f-b8b77d8beae6",
       "multivariate_feature_state_values": []
     },
-    {
-      "django_id": 1212320,
-      "enabled": false,
-      "feature": {
-        "id": 192793,
-        "name": "code_references_ui_stats",
-        "type": "STANDARD"
-      },
-      "feature_segment": null,
-      "feature_state_value": null,
-      "featurestate_uuid": "f976df2f-2341-4623-8425-d6eda23a2ebc",
-      "multivariate_feature_state_values": []
-    },
     {
       "django_id": 1229327,
       "enabled": false,

api/projects/code_references/constants.py

@@ -1,5 +1,2 @@
-# TODO: Implement history cleanup?
-FEATURE_FLAG_CODE_REFERENCES_RETENTION_DAYS = 30
-
 # Linux maximum file path length, as per limits.h/PATH_MAX
 MAX_FILE_PATH_LENGTH = 4096

api/projects/code_references/migrations/0003_introduce_per_feature_scanned_references.py

@@ -0,0 +1,215 @@
+import hashlib
+import json
+from itertools import groupby
+from operator import attrgetter
+from typing import TypedDict
+
+import django.db.models.deletion
+from django.apps.registry import Apps
+from django.db import migrations, models
+from django.db.models import Max
+
+
+class LegacyCodeReference(TypedDict):
+    feature_name: str
+    file_path: str
+    line_number: int
+
+
+class StoredCodeReference(TypedDict):
+    file_path: str
+    line_number: int
+
+
+def _hash_references(references: list[StoredCodeReference]) -> str:
+    return hashlib.md5(json.dumps(references, sort_keys=True).encode()).hexdigest()
+
+
+def migrate_scans_forward(apps: Apps, _: object) -> None:
+    """Split each legacy scan into new cardinality (per-repository and per-feature)"""
+
+    LegacyScan = apps.get_model("code_references", "FeatureFlagCodeReferencesScan")
+    PerFeatureScan = apps.get_model("code_references", "ScannedCodeReferences")
+    Repository = apps.get_model("code_references", "VCSRepository")
+    Feature = apps.get_model("features", "Feature")
+    PerFeatureScan._meta.get_field("created_at").auto_now_add = False
+
+    legacy_scans_summaries = LegacyScan.objects.values(
+        "project_id",
+        "repository_url",
+        "vcs_provider",
+    ).annotate(last_scanned_at=Max("created_at"))
+
+    repositories = {
+        (summary["project_id"], summary["repository_url"]): Repository.objects.create(
+            project_id=summary["project_id"],
+            url=summary["repository_url"],
+            vcs_provider=summary["vcs_provider"],
+            last_scanned_at=summary["last_scanned_at"],
+        )
+        for summary in legacy_scans_summaries
+    }
+
+    # Oldest-first per project so the newest scan wins on hash collisions
+    legacy_scans = LegacyScan.objects.order_by("project_id", "created_at").iterator()
+    grouped_scans = groupby(legacy_scans, key=attrgetter("project_id"))
+    for project_id, project_scans in grouped_scans:
+        features = {
+            (feature.project_id, feature.name): feature
+            for feature in Feature.objects.filter(
+                project_id=project_id,
+                deleted_at__isnull=True,  # Historical models drop SoftDeleteManager
+            )
+        }
+        for legacy_scan in project_scans:
+            repository_url = legacy_scan.repository_url
+            repository = repositories[project_id, repository_url]
+
+            references_by_feature: dict[str, list[StoredCodeReference]] = {}
+            for reference in legacy_scan.code_references:
+                feature_name = reference["feature_name"]
+                references_by_feature.setdefault(feature_name, []).append(
+                    StoredCodeReference(
+                        file_path=reference["file_path"],
+                        line_number=reference["line_number"],
+                    )
+                )
+
+            for feature_name, references in references_by_feature.items():
+                if not (feature := features.get((project_id, feature_name))):
+                    continue
+                PerFeatureScan.objects.update_or_create(
+                    feature=feature,
+                    repository=repository,
+                    code_references_hash=_hash_references(references),
+                    defaults={
+                        "revision": legacy_scan.revision,
+                        "code_references": references,
+                        "created_at": legacy_scan.created_at,
+                    },
+                )
+
+
+def migrate_scans_backward(apps: Apps, _: object) -> None:
+    """Mirror each per-feature row back into the legacy single-table layout."""
+    LegacyScan = apps.get_model("code_references", "FeatureFlagCodeReferencesScan")
+    PerFeatureScan = apps.get_model("code_references", "ScannedCodeReferences")
+    LegacyScan._meta.get_field("created_at").auto_now_add = False
+
+    per_feature_scans = PerFeatureScan.objects.select_related(
+        "repository",
+        "feature",
+    ).iterator(chunk_size=200)
+
+    for per_feature_scan in per_feature_scans:
+        repository = per_feature_scan.repository
+        feature_name = per_feature_scan.feature.name
+        LegacyScan.objects.create(
+            project_id=repository.project_id,
+            repository_url=repository.url,
+            vcs_provider=repository.vcs_provider,
+            revision=per_feature_scan.revision,
+            code_references=[
+                {"feature_name": feature_name, **reference}
+                for reference in per_feature_scan.code_references
+            ],
+            created_at=per_feature_scan.created_at,
+        )
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("code_references", "0002_add_project_repo_created_index"),
+        ("features", "0066_constrain_feature_type"),
+        ("projects", "0029_bump_default_project_limits"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="VCSRepository",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("url", models.URLField()),
+                (
+                    "vcs_provider",
+                    models.CharField(
+                        choices=[("github", "GitHub")],
+                        max_length=50,
+                    ),
+                ),
+                ("last_scanned_at", models.DateTimeField(null=True)),
+                (
+                    "project",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="vcs_repositories",
+                        to="projects.project",
+                    ),
+                ),
+            ],
+        ),
+        migrations.AddConstraint(
+            model_name="vcsrepository",
+            constraint=models.UniqueConstraint(
+                fields=("project", "url"),
+                name="unique_vcs_repository",
+            ),
+        ),
+        migrations.CreateModel(
+            name="ScannedCodeReferences",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("revision", models.CharField(max_length=100)),
+                ("code_references", models.JSONField(default=list)),
+                ("code_references_hash", models.CharField(max_length=32)),
+                (
+                    "feature",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="scanned_code_references",
+                        to="features.feature",
+                    ),
+                ),
+                (
+                    "repository",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="scanned_code_references",
+                        to="code_references.vcsrepository",
+                    ),
+                ),
+            ],
+        ),
+        migrations.AddConstraint(
+            model_name="scannedcodereferences",
+            constraint=models.UniqueConstraint(
+                fields=("feature", "repository", "code_references_hash"),
+                name="unique_scanned_code_references",
+            ),
+        ),
+        migrations.RunPython(
+            code=migrate_scans_forward,
+            reverse_code=migrate_scans_backward,
+        ),
+        migrations.DeleteModel(
+            name="FeatureFlagCodeReferencesScan",
+        ),
+    ]

api/projects/code_references/models.py

@@ -1,37 +1,69 @@
 from django.db import models
 
-from projects.code_references.types import JSONCodeReference, VCSProvider
+from projects.code_references.types import StoredCodeReference, VCSProvider
 
 
-class FeatureFlagCodeReferencesScan(models.Model):
+class VCSRepository(models.Model):
     """
-    A scan of feature flag code references in a repository
+    A VCS repository that is scanned for feature flag code references
     """
 
+    created_at = models.DateTimeField(auto_now_add=True)
+
     project = models.ForeignKey(
         "projects.Project",
         on_delete=models.CASCADE,
-        related_name="code_references",
+        related_name="vcs_repositories",
     )
 
     # Provider-agnostic URL to the web UI of the repository, e.g. https://github.flagsmith.com/backend/
-    repository_url = models.URLField()
+    url = models.URLField()
 
     vcs_provider = models.CharField(
         max_length=50,
         choices=VCSProvider.choices,
-        default=VCSProvider.GITHUB,  # TODO: Remove when adding other providers
     )
+
+    last_scanned_at = models.DateTimeField(null=True)
+
+    class Meta:
+        constraints = [
+            models.UniqueConstraint(
+                fields=["project", "url"],
+                name="unique_vcs_repository",
+            ),
+        ]
+
+
+class ScannedCodeReferences(models.Model):
+    """
+    A list of code references for a feature scanned from a VCS repository
+    """
+
+    created_at = models.DateTimeField(auto_now_add=True)
+
+    feature = models.ForeignKey(
+        "features.Feature",
+        on_delete=models.CASCADE,
+        related_name="scanned_code_references",
+    )
+
+    repository = models.ForeignKey(
+        VCSRepository,
+        on_delete=models.CASCADE,
+        related_name="scanned_code_references",
+    )
+
     revision = models.CharField(max_length=100)
-    code_references = models.JSONField[list[JSONCodeReference]](default=list)
 
-    created_at = models.DateTimeField(auto_now_add=True, db_index=True)
+    code_references = models.JSONField[list[StoredCodeReference]](default=list)
+
+    code_references_hash = models.CharField(max_length=32)
 
     class Meta:
-        ordering = ["-created_at"]
-        indexes = [
-            models.Index(
-                fields=["project", "repository_url", "-created_at"],
-                name="code_ref_proj_repo_created_idx",
+        constraints = [
+            models.UniqueConstraint(
+                fields=["feature", "repository", "code_references_hash"],
+                name="unique_scanned_code_references",
             ),
         ]