From 3f7c9daf0ba00ffd93d52b9e8eb5a6af1a7fbf6d Mon Sep 17 00:00:00 2001 From: Matthew Elwell Date: Fri, 7 Feb 2025 12:51:53 +0100 Subject: [PATCH] fix: delete identity overrides from environments_v2 table on feature delete (#5080) Co-authored-by: Kim Gustyr --- api/edge_api/identities/tasks.py | 11 ++++ .../dynamodb/wrappers/environment_wrapper.py | 18 +++++ api/features/models.py | 15 +++++ .../unit/features/test_unit_features_views.py | 66 +++++++++++++++++++ 4 files changed, 110 insertions(+) diff --git a/api/edge_api/identities/tasks.py b/api/edge_api/identities/tasks.py index 832c2b232b57..1bb23ad8ea2c 100644 --- a/api/edge_api/identities/tasks.py +++ b/api/edge_api/identities/tasks.py @@ -159,3 +159,14 @@ def update_flagsmith_environments_v2_identity_overrides( identifier=identifier, ) dynamodb_wrapper_v2.update_identity_overrides(identity_override_changeset) + + +@register_task_handler() +def delete_environments_v2_identity_overrides_by_feature(feature_id: int) -> None: + dynamodb_wrapper_v2 = DynamoEnvironmentV2Wrapper() + + feature = Feature.objects.all_with_deleted().get(id=feature_id) + for environment in feature.project.environments.all(): + dynamodb_wrapper_v2.delete_identity_overrides( + environment_id=environment.id, feature_id=feature_id + ) diff --git a/api/environments/dynamodb/wrappers/environment_wrapper.py b/api/environments/dynamodb/wrappers/environment_wrapper.py index 6ccb3e911798..56d90b4f25b3 100644 --- a/api/environments/dynamodb/wrappers/environment_wrapper.py +++ b/api/environments/dynamodb/wrappers/environment_wrapper.py @@ -180,3 +180,21 @@ def delete_environment(self, environment_id: int): ENVIRONMENTS_V2_SORT_KEY: item["document_key"], }, ) + + def delete_identity_overrides(self, environment_id: int, feature_id: int) -> None: + filter_expression = self.get_identity_overrides_key_condition_expression( + environment_id=environment_id, feature_id=feature_id + ) + query_kwargs: "QueryInputRequestTypeDef" = { + "KeyConditionExpression": filter_expression, + "ProjectionExpression": "document_key", + } + + with self.table.batch_writer() as writer: + for item in self.query_iter_all_items(**query_kwargs): + writer.delete_item( + Key={ + ENVIRONMENTS_V2_PARTITION_KEY: str(environment_id), + ENVIRONMENTS_V2_SORT_KEY: item["document_key"], + } + ) diff --git a/api/features/models.py b/api/features/models.py index 9eef028114e4..d135b27f92d7 100644 --- a/api/features/models.py +++ b/api/features/models.py @@ -162,6 +162,21 @@ def create_github_comment(self) -> None: def create_feature_states(self): FeatureState.create_initial_feature_states_for_feature(feature=self) + @hook(AFTER_SAVE) + def delete_identity_overrides(self) -> None: + # Note that we have to use conditional logic on self.deleted_at inside + # the hook method because the django-lifecycle logic for when / was / is_not + # doesn't work due to it relying on a method called initial_value, which + # we already define as a field on the model class. + if self.deleted_at and self.project.enable_dynamo_db: + from edge_api.identities.tasks import ( + delete_environments_v2_identity_overrides_by_feature, + ) + + delete_environments_v2_identity_overrides_by_feature.delay( + kwargs={"feature_id": self.id} + ) + def validate_unique(self, *args, **kwargs): """ Checks unique constraints on the model and raises ``ValidationError`` diff --git a/api/tests/unit/features/test_unit_features_views.py b/api/tests/unit/features/test_unit_features_views.py index 893bf7238e6a..4ff89c149bfa 100644 --- a/api/tests/unit/features/test_unit_features_views.py +++ b/api/tests/unit/features/test_unit_features_views.py @@ -1,4 +1,5 @@ import json +import typing import uuid from datetime import date, datetime, timedelta from unittest import mock @@ -31,6 +32,10 @@ IDENTITY_FEATURE_STATE_UPDATED_MESSAGE, ) from audit.models import AuditLog, RelatedObjectType +from environments.dynamodb import ( + DynamoEnvironmentV2Wrapper, + DynamoIdentityWrapper, +) from environments.identities.models import Identity from environments.models import Environment, EnvironmentAPIKey from environments.permissions.models import UserEnvironmentPermission @@ -50,8 +55,17 @@ WithProjectPermissionsCallable, ) from users.models import FFAdminUser, UserPermissionGroup +from util.mappers import ( + map_engine_feature_state_to_identity_override, + map_engine_identity_to_identity_document, + map_identity_override_to_identity_override_document, + map_identity_to_engine, +) from webhooks.webhooks import WebhookEventType +if typing.TYPE_CHECKING: + from mypy_boto3_dynamodb.service_resource import Table + # patch this function as it's triggering extra threads and causing errors mock.patch("features.signals.trigger_feature_state_change_webhooks").start() @@ -3506,3 +3520,55 @@ def test_filter_features_with_owners_and_group_owners_together( assert len(response.data["results"]) == 2 assert response.data["results"][0]["id"] == feature.id assert response.data["results"][1]["id"] == feature2.id + + +def test_delete_feature_deletes_any_related_identity_overrides( + flagsmith_environments_v2_table: "Table", + flagsmith_identities_table: "Table", + dynamodb_identity_wrapper: DynamoIdentityWrapper, + dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper, + environment: Environment, + feature: Feature, + identity_featurestate: FeatureState, + identity: Identity, + admin_client_new: APIClient, +) -> 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) + ) + dynamodb_wrapper_v2.write_environments(environments=[environment]) + flagsmith_environments_v2_table.put_item( + Item=map_identity_override_to_identity_override_document( + map_engine_feature_state_to_identity_override( + feature_state=engine_identity.identity_features[0], + identity_uuid=str(engine_identity.identity_uuid), + identifier=engine_identity.identifier, + environment_api_key=environment.api_key, + environment_id=environment.id, + ), + ) + ) + + feature.project.enable_dynamo_db = True + feature.project.save() + + delete_feature_url = reverse( + "api-v1:projects:project-features-detail", args=[feature.project_id, feature.id] + ) + + # When + delete_feature_response = admin_client_new.delete(delete_feature_url) + + # Then + assert delete_feature_response.status_code == status.HTTP_204_NO_CONTENT + + assert ( + flagsmith_environments_v2_table.query( + KeyConditionExpression=dynamodb_wrapper_v2.get_identity_overrides_key_condition_expression( + environment_id=environment.id, feature_id=feature.id + ) + )["Count"] + == 0 + )