feat: async the logic for cloning feature states into a cloned enviro…

…nment (#4005)

api/environments/migrations/0036_add_is_creating_field.py

@@ -0,0 +1,29 @@
+# Generated by Django 4.2.15 on 2024-10-28 16:18
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("environments", "0035_add_use_identity_overrides_in_local_eval"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="environment",
+            name="is_creating",
+            field=models.BooleanField(
+                default=False,
+                help_text="Attribute used to indicate when an environment is still being created (via clone for example)",
+            ),
+        ),
+        migrations.AddField(
+            model_name="historicalenvironment",
+            name="is_creating",
+            field=models.BooleanField(
+                default=False,
+                help_text="Attribute used to indicate when an environment is still being created (via clone for example)",
+            ),
+        ),
+    ]

api/environments/models.py

@@ -41,7 +41,6 @@
 from environments.managers import EnvironmentManager
 from features.models import Feature, FeatureSegment, FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
-from features.versioning.models import EnvironmentFeatureVersion
 from metadata.models import Metadata
 from projects.models import Project
 from segments.models import Segment
@@ -136,6 +135,12 @@ class Environment(
         default=True,
         help_text="When enabled, identity overrides will be included in the environment document",
     )
+
+    is_creating = models.BooleanField(
+        default=False,
+        help_text="Attribute used to indicate when an environment is still being created (via clone for example)",
+    )
+
     objects = EnvironmentManager()
 
     class Meta:
@@ -163,7 +168,9 @@ def __str__(self):
     def natural_key(self):
         return (self.api_key,)
 
-    def clone(self, name: str, api_key: str = None) -> "Environment":
+    def clone(
+        self, name: str, api_key: str = None, clone_feature_states_async: bool = False
+    ) -> "Environment":
         """
         Creates a clone of the environment, related objects and returns the
         cloned object after saving it to the database.
@@ -173,36 +180,17 @@ def clone(self, name: str, api_key: str = None) -> "Environment":
         clone.id = None
         clone.name = name
         clone.api_key = api_key if api_key else create_hash()
+        clone.is_creating = True
         clone.save()
 
-        # Since identities are closely tied to the environment
-        # it does not make much sense to clone them, hence
-        # only clone feature states without identities
-        queryset = self.feature_states.filter(identity=None)
+        from environments.tasks import clone_environment_feature_states
 
-        if self.use_v2_feature_versioning:
-            # Grab the latest feature versions from the source environment.
-            latest_environment_feature_versions = (
-                EnvironmentFeatureVersion.objects.get_latest_versions_as_queryset(
-                    environment_id=self.id
-                )
-            )
-
-            # Create a dictionary holding the environment feature versions (unique per feature)
-            # to use in the cloned environment.
-            clone_environment_feature_versions = {
-                efv.feature_id: efv.clone_to_environment(environment=clone)
-                for efv in latest_environment_feature_versions
-            }
+        kwargs = {"source_environment_id": self.id, "clone_environment_id": clone.id}
 
-            for feature_state in queryset.filter(
-                environment_feature_version__in=latest_environment_feature_versions
-            ):
-                clone_efv = clone_environment_feature_versions[feature_state.feature_id]
-                feature_state.clone(clone, environment_feature_version=clone_efv)
+        if clone_feature_states_async:
+            clone_environment_feature_states.delay(kwargs=kwargs)
         else:
-            for feature_state in queryset:
-                feature_state.clone(clone, live_from=feature_state.live_from)
+            clone_environment_feature_states(**kwargs)
 
         return clone
 

api/environments/serializers.py

@@ -30,6 +30,7 @@ class Meta:
             "api_key",
             "minimum_change_request_approvals",
             "allow_client_traits",
+            "is_creating",
         )
 
 
@@ -54,6 +55,7 @@ class Meta:
             "hide_sensitive_data",
             "use_v2_feature_versioning",
             "use_identity_overrides_in_local_eval",
+            "is_creating",
         )
         read_only_fields = ("use_v2_feature_versioning",)
 
@@ -127,15 +129,25 @@ def get_subscription(self) -> typing.Optional[Subscription]:
 
 
 class CloneEnvironmentSerializer(EnvironmentSerializerLight):
+    clone_feature_states_async = serializers.BooleanField(
+        default=False,
+        help_text="If True, the environment will be created immediately, but the feature states "
+        "will be created asynchronously. Environment will have `is_creating: true` until "
+        "this process is completed.",
+    )
+
     class Meta:
         model = Environment
-        fields = ("id", "name", "api_key", "project")
+        fields = ("id", "name", "api_key", "project", "clone_feature_states_async")
         read_only_fields = ("id", "api_key", "project")
 
     def create(self, validated_data):
         name = validated_data.get("name")
         source_env = validated_data.get("source_env")
-        clone = source_env.clone(name)
+        clone_feature_states_async = validated_data.get("clone_feature_states_async")
+        clone = source_env.clone(
+            name, clone_feature_states_async=clone_feature_states_async
+        )
         return clone
 
 

api/environments/tasks.py

@@ -8,6 +8,7 @@
     environment_v2_wrapper,
     environment_wrapper,
 )
+from features.versioning.models import EnvironmentFeatureVersion
 from sse import (
     send_environment_update_message_for_environment,
     send_environment_update_message_for_project,
@@ -51,3 +52,43 @@ def delete_environment_from_dynamo(api_key: str, environment_id: str):
 @register_task_handler()
 def delete_environment(environment_id: int) -> None:
     Environment.objects.get(id=environment_id).delete()
+
+
+@register_task_handler()
+def clone_environment_feature_states(
+    source_environment_id: int, clone_environment_id: int
+) -> None:
+    source = Environment.objects.get(id=source_environment_id)
+    clone = Environment.objects.get(id=clone_environment_id)
+
+    # Since identities are closely tied to the environment
+    # it does not make much sense to clone them, hence
+    # only clone feature states without identities
+    queryset = source.feature_states.filter(identity=None)
+
+    if source.use_v2_feature_versioning:
+        # Grab the latest feature versions from the source environment.
+        latest_environment_feature_versions = (
+            EnvironmentFeatureVersion.objects.get_latest_versions_as_queryset(
+                environment_id=source.id
+            )
+        )
+
+        # Create a dictionary holding the environment feature versions (unique per feature)
+        # to use in the cloned environment.
+        clone_environment_feature_versions = {
+            efv.feature_id: efv.clone_to_environment(environment=clone)
+            for efv in latest_environment_feature_versions
+        }
+
+        for feature_state in queryset.filter(
+            environment_feature_version__in=latest_environment_feature_versions
+        ):
+            clone_efv = clone_environment_feature_versions[feature_state.feature_id]
+            feature_state.clone(clone, environment_feature_version=clone_efv)
+    else:
+        for feature_state in queryset:
+            feature_state.clone(clone, live_from=feature_state.live_from)
+
+    clone.is_creating = False
+    clone.save()

api/tests/unit/environments/test_unit_environments_models.py

@@ -145,6 +145,9 @@ def test_environment_clone_clones_the_feature_states(
     # Then
     assert clone.feature_states.first().enabled is True
 
+    clone.refresh_from_db()
+    assert clone.is_creating is False
+
 
 def test_environment_clone_clones_multivariate_feature_state_values(
     environment: Environment,
@@ -994,3 +997,27 @@ def test_clone_environment_v2_versioning(
         cloned_environment_flags.get(feature_segment__segment=segment).enabled
         is expected_segment_fs_enabled_value
     )
+
+
+def test_environment_clone_async(
+    environment: Environment, mocker: MockerFixture
+) -> None:
+    # Given
+    mocked_clone_environment_fs_task = mocker.patch(
+        "environments.tasks.clone_environment_feature_states"
+    )
+
+    # When
+    cloned_environment = environment.clone(
+        name="Cloned environment", clone_feature_states_async=True
+    )
+
+    # Then
+    assert cloned_environment.id != environment.id
+    assert cloned_environment.is_creating is True
+    mocked_clone_environment_fs_task.delay.assert_called_once_with(
+        kwargs={
+            "source_environment_id": environment.id,
+            "clone_environment_id": cloned_environment.id,
+        }
+    )