feat: Heal all identities with blank traits

api/edge_api/management/commands/ensure_identity_traits_blanks.py

@@ -0,0 +1,28 @@
+import logging
+from typing import Any
+
+from django.core.management import BaseCommand
+
+from environments.dynamodb import DynamoIdentityWrapper
+
+identity_wrapper = DynamoIdentityWrapper()
+
+
+logger = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    def handle(self, *args: Any, **options: Any) -> None:
+        for identity_document in identity_wrapper.query_get_all_items():
+            should_write_identity_document = False
+            if identity_traits_data := identity_document.get("identity_traits"):
+                for trait_data in identity_traits_data:
+                    if "trait_value" not in trait_data:
+                        should_write_identity_document = True
+                        trait_data["trait_value"] = ""
+            if should_write_identity_document:
+                identity_wrapper.put_item(identity_document)
+                self.stdout.write(
+                    f"Fixed identity id={identity_document['identity_uuid']}",
+                )
+        self.stdout.write(self.style.SUCCESS("Finished."))

api/environments/dynamodb/wrappers/base.py

@@ -1,4 +1,5 @@
 import typing
+from functools import partial
 
 import boto3
 from botocore.config import Config
@@ -33,13 +34,18 @@ def is_enabled(self) -> bool:
         return self.table is not None
 
     def query_get_all_items(self, **kwargs: dict) -> typing.Generator[dict, None, None]:
+        if kwargs:
+            response_getter = partial(self.table.query, **kwargs)
+        else:
+            response_getter = partial(self.table.scan)
+
         while True:
-            query_response = self.table.query(**kwargs)
+            query_response = response_getter()
             for item in query_response["Items"]:
                 yield item
 
             last_evaluated_key = query_response.get("LastEvaluatedKey")
             if not last_evaluated_key:
                 break
 
-            kwargs["ExclusiveStartKey"] = last_evaluated_key
+            response_getter.keywords["ExclusiveStartKey"] = last_evaluated_key

api/tests/unit/edge_api/test_unit_edge_api_commands.py

@@ -1,8 +1,16 @@
+import typing
+
 from django.core.management import call_command
 from pytest_mock import MockerFixture
 
+from edge_api.management.commands.ensure_identity_traits_blanks import (
+    identity_wrapper,
+)
 from projects.models import EdgeV2MigrationStatus, Project
 
+if typing.TYPE_CHECKING:
+    from mypy_boto3_dynamodb.service_resource import Table
+
 
 def test_migrate_to_edge_v2__new_projects__dont_migrate(
     mocker: MockerFixture, project: Project
@@ -76,3 +84,63 @@ def test_migrate_to_edge_v2__core_projects__dont_migrate(
     # Then
     # unmigrated Core projects were not migrated
     migrate_project_environments_to_v2_mock.assert_not_called()
+
+
+def test_ensure_identity_traits_blanks__calls_expected(
+    flagsmith_identities_table: "Table",
+    mocker: MockerFixture,
+) -> None:
+    # Given
+    environment_api_key = "test"
+    identity_without_traits = {
+        "composite_key": f"{environment_api_key}_identity_without_traits",
+        "environment_api_key": environment_api_key,
+        "identifier": "identity_without_traits",
+        "identity_uuid": "8208c268-e286-4bff-848a-e4b97032fca9",
+    }
+    identity_with_correct_traits = {
+        "composite_key": f"{environment_api_key}_identity_with_correct_traits",
+        "identifier": "identity_with_correct_traits",
+        "environment_api_key": environment_api_key,
+        "identity_uuid": "1a47c1e2-4a9d-4f45-840e-a4cf1a23329e",
+        "identity_traits": [{"trait_key": "key", "trait_value": "value"}],
+    }
+    identity_with_skipped_blank_trait_value = {
+        "composite_key": f"{environment_api_key}_identity_with_skipped_blank_trait_value",
+        "identifier": "identity_with_skipped_blank_trait_value",
+        "environment_api_key": environment_api_key,
+        "identity_uuid": "33e11400-3a34-4b09-9541-3c99e9bf713a",
+        "identity_traits": [
+            {"trait_key": "key", "trait_value": "value"},
+            {"trait_key": "blank"},
+        ],
+    }
+    fixed_identity_with_skipped_blank_trait_value = {
+        **identity_with_skipped_blank_trait_value,
+        "identity_traits": [
+            {"trait_key": "key", "trait_value": "value"},
+            {"trait_key": "blank", "trait_value": ""},
+        ],
+    }
+
+    flagsmith_identities_table.put_item(Item=identity_without_traits)
+    flagsmith_identities_table.put_item(Item=identity_with_correct_traits)
+    flagsmith_identities_table.put_item(Item=identity_with_skipped_blank_trait_value)
+
+    identity_wrapper_put_item_mock = mocker.patch(
+        "edge_api.management.commands.ensure_identity_traits_blanks.identity_wrapper.put_item",
+        side_effect=identity_wrapper.put_item,
+    )
+
+    # When
+    call_command("ensure_identity_traits_blanks")
+
+    # Then
+    assert flagsmith_identities_table.scan()["Items"] == [
+        identity_without_traits,
+        identity_with_correct_traits,
+        fixed_identity_with_skipped_blank_trait_value,
+    ]
+    identity_wrapper_put_item_mock.assert_called_once_with(
+        fixed_identity_with_skipped_blank_trait_value,
+    )