perf: Denormalise feature names on Code Reference scans and add GIN index

api/projects/code_references/db_helpers.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+from django.db.backends.base.base import BaseDatabaseWrapper
+from django.db.models import BooleanField, Func
+from django.db.models.sql.compiler import SQLCompiler
+
+
+class ArrayContains(Func):
+    """Generates: array_col @> ARRAY[value]::text[]
+
+    Used to check whether a text array column contains a single expression
+    value, in a form that PostgreSQL can satisfy with a GIN index. The
+    standard ArrayField __contains lookup only accepts concrete Python values,
+    not ORM expressions such as F() or OuterRef(), hence this helper.
+    """
+
+    output_field = BooleanField()
+
+    def as_sql(
+        self,
+        compiler: SQLCompiler,
+        connection: BaseDatabaseWrapper,
+        *_: Any,
+        **extra_context: Any,
+    ) -> tuple[str, list[str | int] | tuple[str | int, ...] | tuple[()]]:
+        array_expr, value_expr = self.source_expressions
+        array_sql, array_params = compiler.compile(array_expr)
+        value_sql, value_params = compiler.compile(value_expr)
+        return f"{array_sql} @> ARRAY[{value_sql}]::text[]", [
+            *array_params,
+            *value_params,
+        ]

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

@@ -0,0 +1,35 @@
+from django.apps.registry import Apps
+from django.contrib.postgres.fields import ArrayField
+from django.db import migrations, models
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+
+
+def backfill_feature_names(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
+    FeatureFlagCodeReferencesScan = apps.get_model(
+        "code_references", "FeatureFlagCodeReferencesScan"
+    )
+    scans = list(FeatureFlagCodeReferencesScan.objects.all())
+    for scan in scans:
+        scan.feature_names = sorted(
+            {ref["feature_name"] for ref in scan.code_references}
+        )
+    FeatureFlagCodeReferencesScan.objects.bulk_update(scans, ["feature_names"])
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("code_references", "0002_add_project_repo_created_index"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="featureflagcodereferencesscan",
+            field=ArrayField(models.TextField(), default=list),
+            name="feature_names",
+        ),
+        migrations.RunPython(
+            backfill_feature_names,
+            reverse_code=migrations.RunPython.noop,
+        ),
+    ]

api/projects/code_references/migrations/0004_add_feature_names_gin_index.py

@@ -0,0 +1,34 @@
+from django.contrib.postgres.indexes import GinIndex
+from django.db import migrations
+
+from core.migration_helpers import PostgresOnlyRunSQL
+
+
+class Migration(migrations.Migration):
+
+    atomic = False
+
+    dependencies = [
+        ("code_references", "0003_add_feature_names"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            state_operations=[
+                migrations.AddIndex(
+                    model_name="featureflagcodereferencesscan",
+                    index=GinIndex(
+                        fields=["feature_names"],
+                        name="code_refs_feat_names_gin_idx",
+                    ),
+                ),
+            ],
+            database_operations=[
+                PostgresOnlyRunSQL(
+                    'CREATE INDEX CONCURRENTLY IF NOT EXISTS "code_refs_feat_names_gin_idx" '
+                    'ON "code_references_featureflagcodereferencesscan" USING gin ("feature_names");',
+                    reverse_sql='DROP INDEX CONCURRENTLY IF EXISTS "code_refs_feat_names_gin_idx"',
+                ),
+            ],
+        ),
+    ]

api/projects/code_references/models.py

@@ -1,31 +1,42 @@
+from django.contrib.postgres.fields import ArrayField
+from django.contrib.postgres.indexes import GinIndex
 from django.db import models
+from django_lifecycle import (  # type: ignore[import-untyped]
+    BEFORE_SAVE,
+    LifecycleModel,
+    hook,
+)
 
 from projects.code_references.types import JSONCodeReference, VCSProvider
 
 
-class FeatureFlagCodeReferencesScan(models.Model):
+class FeatureFlagCodeReferencesScan(LifecycleModel):  # type: ignore[misc]
     """
     A scan of feature flag code references in a repository
     """
 
-    project = models.ForeignKey(
+    project = models.ForeignKey(  # type: ignore[var-annotated]
         "projects.Project",
         on_delete=models.CASCADE,
         related_name="code_references",
     )
 
     # Provider-agnostic URL to the web UI of the repository, e.g. https://github.flagsmith.com/backend/
-    repository_url = models.URLField()
+    repository_url = models.URLField()  # type: ignore[var-annotated]
 
-    vcs_provider = models.CharField(
+    vcs_provider = models.CharField(  # type: ignore[var-annotated]
         max_length=50,
         choices=VCSProvider.choices,
         default=VCSProvider.GITHUB,  # TODO: Remove when adding other providers
     )
-    revision = models.CharField(max_length=100)
+    revision = models.CharField(max_length=100)  # type: ignore[var-annotated]
     code_references = models.JSONField[list[JSONCodeReference]](default=list)
 
-    created_at = models.DateTimeField(auto_now_add=True, db_index=True)
+    created_at = models.DateTimeField(auto_now_add=True, db_index=True)  # type: ignore[var-annotated]
+
+    # Denormalised from code_references for efficient indexed lookups.
+    # Populated automatically before save and kept in sorted order.
+    feature_names = ArrayField(models.TextField(), default=list)  # type: ignore[var-annotated]
 
     class Meta:
         ordering = ["-created_at"]
@@ -34,4 +45,14 @@ class Meta:
                 fields=["project", "repository_url", "-created_at"],
                 name="code_ref_proj_repo_created_idx",
             ),
+            GinIndex(
+                fields=["feature_names"],
+                name="code_refs_feat_names_gin_idx",
+            ),
         ]
+
+    @hook(BEFORE_SAVE)  # type: ignore[misc]
+    def populate_feature_names(self) -> None:
+        self.feature_names = sorted(
+            {ref["feature_name"] for ref in self.code_references}
+        )

api/projects/code_references/services.py

@@ -4,7 +4,6 @@
 from django.contrib.postgres.expressions import ArraySubquery
 from django.contrib.postgres.fields import ArrayField
 from django.db.models import (
-    BooleanField,
     F,
     Func,
     JSONField,
@@ -20,6 +19,7 @@
 from projects.code_references.constants import (
     FEATURE_FLAG_CODE_REFERENCES_RETENTION_DAYS,
 )
+from projects.code_references.db_helpers import ArrayContains
 from projects.code_references.models import FeatureFlagCodeReferencesScan
 from projects.code_references.types import (
     CodeReference,
@@ -53,13 +53,7 @@ def annotate_feature_queryset_with_code_references_summary(
     last_feature_found_at = (
         FeatureFlagCodeReferencesScan.objects.annotate(
             feature_name=OuterRef("feature_name"),
-            contains_feature_name=Func(
-                F("code_references"),
-                Value("$[*] ? (@.feature_name == $feature_name)"),
-                JSONObject(feature_name=F("feature_name")),
-                function="jsonb_path_exists",
-                output_field=BooleanField(),
-            ),
+            contains_feature_name=ArrayContains(F("feature_names"), F("feature_name")),
         )
         .filter(
             project=OuterRef("project_id"),
@@ -122,7 +116,7 @@ def get_code_references_for_feature_flag(
             project=feature.project,
             created_at__gte=timezone.now() - history_delta,
             repository_url=OuterRef("repository_url"),
-            code_references__contains=[{"feature_name": feature.name}],
+            feature_names__contains=[feature.name],
         )
         .values("created_at")
         .order_by("-created_at")[:1]

api/tests/unit/projects/code_references/test_unit_projects_code_references_db_helpers.py

@@ -0,0 +1,86 @@
+from typing import Iterator
+
+import pytest
+from django.contrib.postgres.fields import ArrayField
+from django.db import connection, models
+from django.db.models import F, Value
+from django.test.utils import isolate_apps
+
+from projects.code_references.db_helpers import ArrayContains
+
+
+@pytest.fixture()
+def names_model(db: None) -> Iterator[type[models.Model]]:
+    with isolate_apps("projects.code_references"):
+
+        class NamesModel(models.Model):
+            names = ArrayField(models.TextField(), default=list)
+
+            class Meta:
+                app_label = "projects.code_references"
+
+        with connection.schema_editor() as editor:
+            editor.create_model(NamesModel)
+
+        yield NamesModel
+
+
+def test_ArrayContains__matches_when_value_present_in_array(
+    names_model: models.Model,
+) -> None:
+    # Given
+    names_model.objects.create(names=["john", "esme"])  # type: ignore[attr-defined]
+
+    # When
+    result = names_model.objects.annotate(  # type: ignore[attr-defined]
+        has_name=ArrayContains(F("names"), Value("esme")),
+    ).get()
+
+    # Then
+    assert result.has_name is True
+
+
+def test_ArrayContains__does_not_match_when_value_absent_from_array(
+    names_model: models.Model,
+) -> None:
+    # Given
+    names_model.objects.create(names=["john"])  # type: ignore[attr-defined]
+
+    # When
+    result = names_model.objects.annotate(  # type: ignore[attr-defined]
+        has_name=ArrayContains(F("names"), Value("lisa")),
+    ).get()
+
+    # Then
+    assert result.has_name is False
+
+
+def test_ArrayContains__filters_queryset_correctly(
+    names_model: models.Model,
+) -> None:
+    # Given
+    matching = names_model.objects.create(names=["john", "esme"])  # type: ignore[attr-defined]
+    names_model.objects.create(names=["lisa"])  # type: ignore[attr-defined]
+
+    # When
+    results = names_model.objects.annotate(  # type: ignore[attr-defined]
+        has_name=ArrayContains(F("names"), Value("john")),
+    ).filter(has_name=True)
+
+    # Then
+    assert list(results) == [matching]
+
+
+def test_ArrayContains__does_not_match_empty_array(
+    names_model: models.Model,
+) -> None:
+    # Given
+    names_model.objects.create(names=[])  # type: ignore[attr-defined]
+
+    # When
+    result = names_model.objects.annotate(  # type: ignore[attr-defined]
+        has_name=ArrayContains(F("names"), Value("kiefer")),
+    ).get()
+
+    # Then
+    assert result.has_name is False

api/tests/unit/projects/code_references/test_unit_projects_code_references_migrations.py

@@ -0,0 +1,54 @@
+from django_test_migrations.migrator import Migrator
+
+
+def test_backfill_feature_names(migrator: Migrator) -> None:
+    # Given
+    old_state = migrator.apply_initial_migration(
+        ("code_references", "0002_add_project_repo_created_index")
+    )
+
+    Organisation = old_state.apps.get_model("organisations", "Organisation")
+    Project = old_state.apps.get_model("projects", "Project")
+    FeatureFlagCodeReferencesScan = old_state.apps.get_model(
+        "code_references", "FeatureFlagCodeReferencesScan"
+    )
+
+    organisation = Organisation.objects.create(name="Test Organisation")
+    project = Project.objects.create(name="Test Project", organisation=organisation)
+
+    # A scan with references to multiple features (unsorted) including a duplicate
+    scan_with_references = FeatureFlagCodeReferencesScan.objects.create(
+        project=project,
+        repository_url="https://github.com/example/repo",
+        revision="abc123",
+        code_references=[
+            {"feature_name": "zebra_flag", "file_path": "foo.py", "line_number": 1},
+            {"feature_name": "alpha_flag", "file_path": "bar.py", "line_number": 2},
+            {"feature_name": "zebra_flag", "file_path": "baz.py", "line_number": 3},
+        ],
+    )
+
+    # A scan with no code references
+    scan_with_no_references = FeatureFlagCodeReferencesScan.objects.create(
+        project=project,
+        repository_url="https://github.com/example/repo",
+        revision="def456",
+        code_references=[],
+    )
+
+    # When
+    new_state = migrator.apply_tested_migration(
+        ("code_references", "0003_add_feature_names")
+    )
+    NewScan = new_state.apps.get_model(
+        "code_references", "FeatureFlagCodeReferencesScan"
+    )
+
+    # Then
+    # Duplicates are removed and names are sorted
+    updated_scan = NewScan.objects.get(id=scan_with_references.id)
+    assert updated_scan.feature_names == ["alpha_flag", "zebra_flag"]
+
+    # Empty code_references results in empty feature_names
+    updated_scan_no_refs = NewScan.objects.get(id=scan_with_no_references.id)
+    assert updated_scan_no_refs.feature_names == []