fix: delete identity overrides from environments_v2 table on feature delete

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
+        )

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"],
+                    }
+                )

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``

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
+    )