fix: add migration to clean up corrupt data caused by feature versioning

api/features/versioning/migrations/0005_fix_scheduled_fs_data_issue_caused_by_enabling_versioning.py

@@ -0,0 +1,71 @@
+# Generated by Django 4.2.16 on 2024-11-28 13:45
+from django.apps.registry import Apps
+from django.db import migrations
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.models import F
+
+
+def fix_corrupted_feature_states(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
+    feature_state_model_class = apps.get_model("features", "FeatureState")
+    environment_feature_version_model_class = apps.get_model(
+        "feature_versioning", "EnvironmentFeatureVersion"
+    )
+
+    feature_states_to_update = []
+    environment_feature_versions_to_update = []
+
+    for feature_state in feature_state_model_class.objects.filter(
+        # We're looking for feature states that live in environments
+        # that don't have versioning enabled, but have an environment
+        # feature version associated to them. Moreover, those environment
+        # feature versions should have a change request associated with
+        # them. See this PR for further details of the bug we're looking
+        # to clean up after: https://github.com/Flagsmith/flagsmith/pull/4872
+        environment__use_v2_feature_versioning=False,
+        environment_feature_version__isnull=False,
+        environment_feature_version__change_request__isnull=False,
+    ).exclude(
+        # Just to be safe, we exclude feature states for which the version
+        # belongs to the same feature and environment. This will prevent us
+        # from accidentally modifying any feature states that belong to
+        # legitimate environment feature versions but perhaps versioning
+        # has been switched off for the environment.
+        environment_feature_version__feature=F("feature"),
+        environment_feature_version__environment=F("environment")
+    ).select_related("environment_feature_version"):
+        environment_feature_version = feature_state.environment_feature_version
+
+        # 1. move the change request back to the feature state
+        feature_state.change_request = environment_feature_version.change_request
+
+        # 2. remove the change request from the EnvironmentFeatureVersion
+        environment_feature_version.change_request = None
+
+        # 3. remove the environment feature version from the feature state
+        feature_state.environment_feature_version = None
+
+        feature_states_to_update.append(feature_state)
+        environment_feature_versions_to_update.append(environment_feature_version)
+
+    feature_state_model_class.objects.bulk_update(
+        feature_states_to_update,
+        ["environment_feature_version", "change_request"],
+    )
+    environment_feature_version_model_class.objects.bulk_update(
+        environment_feature_versions_to_update,
+        ["change_request"],
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("feature_versioning", "0004_add_version_change_set"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            fix_corrupted_feature_states,
+            reverse_code=migrations.RunPython.noop
+        ),
+    ]

api/tests/unit/features/versioning/test_unit_versioning_migrations.py

@@ -0,0 +1,154 @@
+from datetime import timedelta
+
+from django.utils import timezone
+from django_test_migrations.migrator import Migrator
+
+
+def test_fix_scheduled_fs_data_issue_caused_by_enabling_versioning(
+    migrator: Migrator
+) -> None:
+    # Given
+    now = timezone.now()
+    one_hour_from_now = now + timedelta(hours=1)
+
+    initial_state = migrator.apply_initial_migration(("feature_versioning", "0004_add_version_change_set"))
+
+    organisation_model_class = initial_state.apps.get_model("organisations", "Organisation")
+    project_model_class = initial_state.apps.get_model("projects", "Project")
+    environment_model_class = initial_state.apps.get_model("environments", "Environment")
+    change_request_model_class = initial_state.apps.get_model("workflows_core", "ChangeRequest")
+    environment_feature_version_model_class = initial_state.apps.get_model(
+        "feature_versioning", "EnvironmentFeatureVersion"
+    )
+    feature_state_model_class = initial_state.apps.get_model("features", "FeatureState")
+    feature_model_class = initial_state.apps.get_model("features", "Feature")
+
+    organisation = organisation_model_class.objects.create(name="Test Organisation")
+    project = project_model_class.objects.create(name="Test Project", organisation=organisation)
+    environment_1 = environment_model_class.objects.create(name="Environment 1", project=project)
+    feature = feature_model_class.objects.create(name="test_feature", project=project)
+
+    # First, let's create some regular data that should be left untouched for a
+    # non-versioned environment
+    # Note: because migrations don't trigger the signals, we need to manually create
+    # the default state
+    environment_1_default_feature_state = feature_state_model_class.objects.create(
+        feature=feature,
+        environment=environment_1,
+        version=1,
+    )
+    environment_1_live_cr = change_request_model_class.objects.create(
+        title="Live CR for Environment 1",
+        environment=environment_1,
+    )
+    environment_1_live_feature_state = feature_state_model_class.objects.create(
+        feature=feature,
+        environment=environment_1,
+        enabled=True,
+        live_from=one_hour_from_now,
+        version=2,
+        change_request=environment_1_live_cr,
+    )
+
+    # ... and a versioned environment
+    versioned_environment = environment_model_class.objects.create(
+        name="Versioned Environment",
+        project=project,
+        use_v2_feature_versioning=True,
+    )
+    versioned_environment_version_1 = environment_feature_version_model_class.objects.create(
+        feature=feature,
+        environment=versioned_environment,
+    )
+    versioned_environment_default_feature_state = feature_state_model_class.objects.create(
+        feature=feature,
+        environment=versioned_environment,
+        environment_feature_version=versioned_environment_version_1,
+    )
+    versioned_environment_cr = change_request_model_class.objects.create(
+        environment=versioned_environment,
+        title="Versioned Environment CR",
+    )
+    versioned_environment_cr_version = environment_feature_version_model_class.objects.create(
+        feature=feature,
+        environment=versioned_environment,
+        change_request=versioned_environment_cr,
+    )
+    versioned_environment_change_request_feature_state = feature_state_model_class.objects.create(
+        feature=feature,
+        environment=versioned_environment,
+        environment_feature_version=versioned_environment_cr_version,
+    )
+
+    # Now, let's create the corrupt data that we want to correct
+    environment_2 = environment_model_class.objects.create(name="Environment 2", project=project)
+    # Note: because migrations don't trigger the signals, we need to manually create
+    # the default state
+    feature_state_model_class.objects.create(
+        feature=feature,
+        environment=environment_2,
+        version=1,
+    )
+    environment_2_scheduled_cr = change_request_model_class.objects.create(
+        title="Environment 2 Scheduled CR",
+        environment=environment_2,
+    )
+    environment_2_scheduled_feature_state = feature_state_model_class.objects.create(
+        feature=feature,
+        environment=environment_2,
+        live_from=one_hour_from_now,
+        version=2,
+        change_request=environment_2_scheduled_cr,
+    )
+
+    # and let's explicitly write out the corruption steps as they would have occurred
+    # when any environment has feature versioning enabled.
+    # 1. The scheduled feature state for environment 2 would be associated with a new
+    #    environment feature version in that environment.
+    new_version = environment_feature_version_model_class.objects.create(
+        environment=versioned_environment,
+        feature=feature,
+    )
+    environment_2_scheduled_feature_state.environment_feature_version = new_version
+
+    # 2. The change request would be removed from the feature state
+    environment_2_scheduled_feature_state.change_request = None
+
+    # 3. The change request would be added to the new version instead
+    new_version.change_request = environment_2_scheduled_cr
+
+    # 4. Let's save the models with the modifications
+    new_version.save()
+    environment_2_scheduled_feature_state.save()
+
+    # When
+    new_state = migrator.apply_tested_migration(
+        (
+            "feature_versioning",
+            "0005_fix_scheduled_fs_data_issue_caused_by_enabling_versioning",
+        )
+    )
+
+    # Then
+    feature_state_model_class = new_state.apps.get_model("features", "FeatureState")
+
+    # Let's check that the regular data is untouched
+    assert feature_state_model_class.objects.get(
+        pk=environment_1_default_feature_state.pk
+    ).environment_feature_version is None
+    assert feature_state_model_class.objects.get(
+        pk=environment_1_live_feature_state.pk
+    ).environment_feature_version is None
+    assert feature_state_model_class.objects.get(
+        pk=versioned_environment_default_feature_state.pk
+    ).environment_feature_version_id == versioned_environment_version_1.pk
+    assert feature_state_model_class.objects.get(
+        pk=versioned_environment_change_request_feature_state.pk
+    ).environment_feature_version_id == versioned_environment_cr_version.pk
+
+    # And let's check the corrupted data issue has been solved
+    new_environment_2_scheduled_feature_state = feature_state_model_class.objects.get(
+        pk=environment_2_scheduled_feature_state.pk
+    )
+    assert new_environment_2_scheduled_feature_state.environment_feature_version is None
+    assert new_environment_2_scheduled_feature_state.change_request_id == environment_2_scheduled_cr.pk