From 954614b5e7b8a1d44245047adb65ec4e9b2274dc Mon Sep 17 00:00:00 2001
From: Kim Gustyr <kim.gustyr@flagsmith.com>
Date: Wed, 7 Aug 2024 11:51:31 +0100
Subject: [PATCH] fix: Identity overrides are not deleted when deleting Edge
 identities

---
 api/edge_api/identities/models.py | 57 ++++++++++++++++++++++---------
 api/edge_api/identities/views.py  |  8 +++--
 2 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/api/edge_api/identities/models.py b/api/edge_api/identities/models.py
index d67bcb5831c8..1e062ecd005e 100644
--- a/api/edge_api/identities/models.py
+++ b/api/edge_api/identities/models.py
@@ -167,15 +167,52 @@ def remove_feature_override(self, feature_state: FeatureStateModel) -> None:
 
     def save(self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None):
         self.dynamo_wrapper.put_item(self.to_document())
-        changes = self._get_changes()
-        if changes["feature_overrides"]:
+        changeset = self._get_changes()
+        self._update_feature_overrides(
+            changeset=changeset,
+            user=user,
+            master_api_key=master_api_key,
+        )
+        self._reset_initial_state()
+
+    def delete(
+        self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None
+    ) -> None:
+        self.dynamo_wrapper.delete_item(self._engine_identity_model.composite_key)
+        self._engine_identity_model.identity_features.clear()
+        changeset = self._get_changes()
+        self._update_feature_overrides(
+            changeset=changeset,
+            user=user,
+            master_api_key=master_api_key,
+        )
+        self._reset_initial_state()
+
+    def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None:
+        identity_feature_names = {
+            fs.feature.name for fs in self._engine_identity_model.identity_features
+        }
+        if not identity_feature_names.issubset(valid_feature_names):
+            self._engine_identity_model.prune_features(list(valid_feature_names))
+            sync_identity_document_features.delay(args=(str(self.identity_uuid),))
+
+    def to_document(self) -> dict:
+        return map_engine_identity_to_identity_document(self._engine_identity_model)
+
+    def _update_feature_overrides(
+        self,
+        changeset: IdentityChangeset,
+        user: FFAdminUser = None,
+        master_api_key: MasterAPIKey = None,
+    ) -> None:
+        if changeset["feature_overrides"]:
             # TODO: would this be simpler if we put a wrapper around FeatureStateModel instead?
             generate_audit_log_records.delay(
                 kwargs={
                     "environment_api_key": self.environment_api_key,
                     "identifier": self.identifier,
                     "user_id": getattr(user, "id", None),
-                    "changes": changes,
+                    "changes": changeset,
                     "identity_uuid": str(self.identity_uuid),
                     "master_api_key_id": getattr(master_api_key, "id", None),
                 }
@@ -183,23 +220,11 @@ def save(self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None):
             update_flagsmith_environments_v2_identity_overrides.delay(
                 kwargs={
                     "environment_api_key": self.environment_api_key,
-                    "changes": changes,
+                    "changes": changeset,
                     "identity_uuid": str(self.identity_uuid),
                     "identifier": self.identifier,
                 }
             )
-        self._reset_initial_state()
-
-    def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None:
-        identity_feature_names = {
-            fs.feature.name for fs in self._engine_identity_model.identity_features
-        }
-        if not identity_feature_names.issubset(valid_feature_names):
-            self._engine_identity_model.prune_features(list(valid_feature_names))
-            sync_identity_document_features.delay(args=(str(self.identity_uuid),))
-
-    def to_document(self) -> dict:
-        return map_engine_identity_to_identity_document(self._engine_identity_model)
 
     def _get_changes(self) -> IdentityChangeset:
         previous_instance = self._initial_state
diff --git a/api/edge_api/identities/views.py b/api/edge_api/identities/views.py
index c8ac3c0e3d50..9065324457c5 100644
--- a/api/edge_api/identities/views.py
+++ b/api/edge_api/identities/views.py
@@ -160,7 +160,11 @@ def get_environment_from_request(self):
         )
 
     def perform_destroy(self, instance):
-        EdgeIdentity.dynamo_wrapper.delete_item(instance["composite_key"])
+        edge_identity = EdgeIdentity.from_identity_document(instance)
+        edge_identity.delete(
+            user=self.request.user,
+            master_api_key=getattr(self.request, "master_api_key", None),
+        )
 
     @swagger_auto_schema(
         responses={200: EdgeIdentityTraitsSerializer(many=True)},
@@ -281,7 +285,7 @@ def list(self, request, *args, **kwargs):
     def perform_destroy(self, instance):
         self.identity.remove_feature_override(instance)
         self.identity.save(
-            user=self.request.user.id,
+            user=self.request.user,
             master_api_key=getattr(self.request, "master_api_key", None),
         )