From 02f5f71f82bae1ec3536cb522fc0b684a2c27605 Mon Sep 17 00:00:00 2001 From: Matthew Elwell Date: Wed, 30 Oct 2024 16:22:45 +0000 Subject: [PATCH] feat: async the logic for cloning feature states into a cloned environment (#4005) --- .../migrations/0036_add_is_creating_field.py | 29 +++++++++++++ api/environments/models.py | 42 +++++++------------ api/environments/serializers.py | 16 ++++++- api/environments/tasks.py | 41 ++++++++++++++++++ .../test_unit_environments_models.py | 27 ++++++++++++ 5 files changed, 126 insertions(+), 29 deletions(-) create mode 100644 api/environments/migrations/0036_add_is_creating_field.py diff --git a/api/environments/migrations/0036_add_is_creating_field.py b/api/environments/migrations/0036_add_is_creating_field.py new file mode 100644 index 000000000000..3d467ad24f5e --- /dev/null +++ b/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)", + ), + ), + ] diff --git a/api/environments/models.py b/api/environments/models.py index f85b6851d877..11bf982fd5a5 100644 --- a/api/environments/models.py +++ b/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 diff --git a/api/environments/serializers.py b/api/environments/serializers.py index dffc9d13b468..36d2bea5ee3a 100644 --- a/api/environments/serializers.py +++ b/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 diff --git a/api/environments/tasks.py b/api/environments/tasks.py index 75f66df5f1d4..5767dd8adbdd 100644 --- a/api/environments/tasks.py +++ b/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() diff --git a/api/tests/unit/environments/test_unit_environments_models.py b/api/tests/unit/environments/test_unit_environments_models.py index 046e7a899bba..82c1ab0fe3b6 100644 --- a/api/tests/unit/environments/test_unit_environments_models.py +++ b/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, + } + )