Add ability to skip identity overrides when migrating to environments_v2

Author: khvn26

api/environments/dynamodb/services.py

@@ -17,7 +17,10 @@
 logger = logging.getLogger(__name__)
 
 
-def migrate_environments_to_v2(project_id: int) -> IdentityOverridesV2Changeset | None:
+def migrate_environments_to_v2(
+    project_id: int,
+    skip_identity_overrides: bool,
+) -> IdentityOverridesV2Changeset | None:
     dynamo_wrapper_v2 = DynamoEnvironmentV2Wrapper()
     identity_wrapper = DynamoIdentityWrapper()
 
@@ -27,19 +30,20 @@ def migrate_environments_to_v2(project_id: int) -> IdentityOverridesV2Changeset
     logger.info("Migrating environments to v2 for project %d", project_id)
 
     environments_to_migrate = Environment.objects.filter(project_id=project_id)
-    identity_overrides_to_migrate = _iter_paginated_overrides(
-        identity_wrapper=identity_wrapper,
-        environments=environments_to_migrate,
-    )
 
-    changeset = IdentityOverridesV2Changeset(
-        to_put=list(identity_overrides_to_migrate), to_delete=[]
-    )
-    logger.info(
-        "Retrieved %d identity overrides to migrate for project %d",
-        len(changeset.to_put),
-        project_id,
-    )
+    changeset = IdentityOverridesV2Changeset(to_put=[], to_delete=[])
+    if not skip_identity_overrides:
+        changeset.to_put = [
+            *_iter_paginated_overrides(
+                identity_wrapper=identity_wrapper,
+                environments=environments_to_migrate,
+            )
+        ]
+        logger.info(
+            "Retrieved %d identity overrides to migrate for project %d",
+            len(changeset.to_put),
+            project_id,
+        )
 
     dynamo_wrapper_v2.write_environments(environments_to_migrate)
     dynamo_wrapper_v2.update_identity_overrides(changeset)

api/projects/tasks.py

@@ -17,7 +17,10 @@ def migrate_project_environments_to_v2(project_id: int) -> None:
 
     with transaction.atomic():
         project = Project.objects.select_for_update().get(id=project_id)
-        if migrate_environments_to_v2(project_id=project_id):
+        if migrate_environments_to_v2(
+            project_id=project_id,
+            skip_identity_overrides=True,
+        ):
             project.identity_overrides_v2_migration_status = (
                 IdentityOverridesV2MigrationStatus.COMPLETE
             )

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

@@ -48,7 +48,10 @@ def test_migrate_environments_to_v2__environment_with_overrides__writes_expected
     )
 
     # When
-    migrate_environments_to_v2(project_id=environment.project_id)
+    migrate_environments_to_v2(
+        project_id=environment.project_id,
+        skip_identity_overrides=False,
+    )
 
     # Then
     results = flagsmith_environments_v2_table.scan()["Items"]
@@ -73,8 +76,41 @@ def test_migrate_environments_to_v2__wrapper_disabled__does_not_write(
     )
 
     # When
-    migrate_environments_to_v2(project_id=mocker.Mock())
+    migrate_environments_to_v2(
+        project_id=mocker.Mock(),
+        skip_identity_overrides=False,
+    )
 
     # Then
     mocked_dynamodb_identity_wrapper.return_value.assert_not_called()
     mocked_dynamodb_v2_wrapper.return_value.assert_not_called()
+
+
+def test_migrate_environments_to_v2__project_skip_identity_overrides_true__only_writes_environment(
+    environment: Environment,
+    identity: Identity,
+    identity_featurestate: FeatureState,
+    dynamodb_identity_wrapper: DynamoIdentityWrapper,
+    dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper,
+    flagsmith_environments_v2_table: Table,
+) -> None:
+    # Given
+    engine_identity = map_identity_to_engine(identity, with_overrides=True)
+    dynamodb_identity_wrapper.put_item(
+        map_engine_identity_to_identity_document(engine_identity)
+    )
+
+    expected_environment_document = map_environment_to_environment_v2_document(
+        environment
+    )
+
+    # When
+    migrate_environments_to_v2(
+        project_id=environment.project_id,
+        skip_identity_overrides=True,
+    )
+
+    # Then
+    results = flagsmith_environments_v2_table.scan()["Items"]
+    assert len(results) == 1
+    assert results[0] == expected_environment_document

api/tests/unit/projects/test_unit_projects_tasks.py

@@ -61,6 +61,7 @@ def test_migrate_project_environments_to_v2__calls_expected(
     project_v2_migration_in_progress.refresh_from_db()
     mocked_migrate_environments_to_v2.assert_called_once_with(
         project_id=project_v2_migration_in_progress.id,
+        skip_identity_overrides=False,
     )
     assert (
         project_v2_migration_in_progress.identity_overrides_v2_migration_status
@@ -91,7 +92,8 @@ def test_migrate_project_environments_to_v2__expected_status_on_error(
 
     # Then
     mocked_migrate_environments_to_v2.assert_called_once_with(
-        project_id=project_v2_migration_in_progress.id
+        project_id=project_v2_migration_in_progress.id,
+        skip_identity_overrides=False,
     )
     assert project_v2_migration_in_progress.identity_overrides_v2_migration_status == (
         IdentityOverridesV2MigrationStatus.IN_PROGRESS