fix: replace OFFSET pagination with keyset pagination in identity migration (#6877)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: gagantrivedi

api/environments/dynamodb/migrator.py

@@ -1,3 +1,5 @@
+from collections.abc import Iterator
+
 from django.db.models import Prefetch
 
 from edge_api.identities.events import send_migration_event
@@ -6,7 +8,6 @@
 from features.models import FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
 from projects.models import Project
-from util.queryset import iterator_with_prefetch
 
 from .types import DynamoProjectMetadata, ProjectIdentityMigrationStatus
 from .wrappers import (
@@ -17,6 +18,53 @@
 
 
 class IdentityMigrator:
+    @staticmethod
+    def iter_identities_in_chunks(
+        project_id: int, chunk_size: int = 2000
+    ) -> Iterator[Identity]:
+        """
+        Yield identities in fixed-size chunks using keyset pagination.
+
+        We don't use Django's built-in QuerySet.iterator() here because
+        it uses server-side cursors (DECLARE/FETCH), which plan the
+        entire result set upfront and can hit statement_timeout on
+        large tables.
+
+        Each chunk issues a WHERE pk > last_pk ORDER BY pk LIMIT N query,
+        which is O(1) via pk index seek regardless of table size.
+        """
+        identities_qs = (
+            Identity.objects.filter(environment__project__id=project_id)
+            .select_related("environment")
+            .prefetch_related(
+                "identity_traits",
+                Prefetch(
+                    "identity_features",
+                    queryset=FeatureState.objects.select_related(
+                        "feature", "feature_state_value"
+                    ),
+                ),
+                Prefetch(
+                    "identity_features__multivariate_feature_state_values",
+                    queryset=MultivariateFeatureStateValue.objects.select_related(
+                        "multivariate_feature_option"
+                    ),
+                ),
+            )
+        )
+        queryset = identities_qs.order_by("pk")
+        last_pk = None
+
+        while True:
+            chunk_qs = (
+                queryset.filter(pk__gt=last_pk) if last_pk is not None else queryset
+            )
+            chunk = list(chunk_qs[:chunk_size])
+            if not chunk:
+                break
+            yield from chunk
+            last_pk = chunk[-1].pk
+
     def __init__(self, project_id):  # type: ignore[no-untyped-def]
         self.project_metadata = DynamoProjectMetadata.get_or_new(project_id)
 
@@ -63,24 +111,5 @@ def migrate(self):  # type: ignore[no-untyped-def]
         identity_wrapper = DynamoIdentityWrapper(
             environment_wrapper=environment_wrapper
         )
-        identities = (
-            Identity.objects.filter(environment__project__id=project_id)
-            .select_related("environment")
-            .prefetch_related(
-                "identity_traits",
-                Prefetch(
-                    "identity_features",
-                    queryset=FeatureState.objects.select_related(
-                        "feature", "feature_state_value"
-                    ),
-                ),
-                Prefetch(
-                    "identity_features__multivariate_feature_state_values",
-                    queryset=MultivariateFeatureStateValue.objects.select_related(
-                        "multivariate_feature_option"
-                    ),
-                ),
-            )
-        )
-        identity_wrapper.write_identities(iterator_with_prefetch(identities))  # type: ignore[no-untyped-call]
+        identity_wrapper.write_identities(self.iter_identities_in_chunks(project_id))
         self.project_metadata.finish_identity_migration()  # type: ignore[no-untyped-call]

api/tests/unit/environments/dynamodb/test_unit_migrator.py

@@ -1,12 +1,15 @@
 from pytest_django.asserts import assertQuerySetEqual as assert_queryset_equal
+from pytest_django.fixtures import DjangoAssertNumQueries
 
 from environments.dynamodb.migrator import IdentityMigrator
 from environments.dynamodb.types import (
     DynamoProjectMetadata,
     ProjectIdentityMigrationStatus,
 )
 from environments.identities.models import Identity
+from environments.identities.traits.models import Trait
 from environments.models import Environment, EnvironmentAPIKey
+from projects.models import Project
 
 
 def test_migrate_calls_internal_methods_with_correct_arguments(  # type: ignore[no-untyped-def]
@@ -48,7 +51,7 @@ def test_migrate_calls_internal_methods_with_correct_arguments(  # type: ignore[
     assert kwargs == {}
 
     assert_queryset_equal(
-        args[0], Identity.objects.filter(environment__project__id=project.id)
+        list(args[0]), Identity.objects.filter(environment__project__id=project.id)
     )
     # and
     args, kwargs = mocked_environment_wrapper.return_value.write_environments.call_args
@@ -168,3 +171,70 @@ def test_get_migration_status_returns_correct_migraion_status_for_in_progress_mi
     # Then
     assert status == ProjectIdentityMigrationStatus.MIGRATION_IN_PROGRESS
     mocked_project_metadata.get_or_new.assert_called_with(project_id)
+
+
+def test_iter_identities_in_chunks__multiple_chunks__yields_all_in_pk_order(
+    project: Project,
+    environment: Environment,
+) -> None:
+    # Given
+    identities = [
+        Identity.objects.create(identifier=f"identity_{i}", environment=environment)
+        for i in range(5)
+    ]
+    for identity in identities:
+        Trait.objects.create(
+            identity=identity,
+            trait_key="key",
+            value_type="unicode",
+            string_value="val",
+        )
+
+    # When
+    result = list(IdentityMigrator.iter_identities_in_chunks(project.id, chunk_size=2))
+
+    # Then
+    expected_pks = sorted(i.pk for i in identities)
+    assert [i.pk for i in result] == expected_pks
+    assert len(result) == 5
+
+
+def test_iter_identities_in_chunks__preserves_prefetch_related(
+    project: Project,
+    environment: Environment,
+    django_assert_num_queries: DjangoAssertNumQueries,
+) -> None:
+    # Given
+    identities = [
+        Identity.objects.create(identifier=f"identity_{i}", environment=environment)
+        for i in range(3)
+    ]
+    for identity in identities:
+        Trait.objects.create(
+            identity=identity,
+            trait_key="key",
+            value_type="unicode",
+            string_value="val",
+        )
+
+    result = list(IdentityMigrator.iter_identities_in_chunks(project.id))
+
+    # When — accessing prefetched relations should cause no additional queries.
+    with django_assert_num_queries(0):
+        for identity in result:
+            list(identity.identity_traits.all())
+
+
+def test_iter_identities_in_chunks__empty_queryset__yields_nothing(
+    project: Project,
+    environment: Environment,
+    django_assert_num_queries: DjangoAssertNumQueries,
+) -> None:
+    # Given — no identities created.
+
+    # When
+    with django_assert_num_queries(1):
+        result = list(IdentityMigrator.iter_identities_in_chunks(project.id))
+
+    # Then
+    assert result == []