feat: Heal all identities with blank traits

api/edge_api/management/commands/ensure_identity_traits_blanks.py

@@ -0,0 +1,63 @@
+import logging
+from typing import Any
+
+from django.core.management import BaseCommand
+
+from environments.dynamodb import DynamoIdentityWrapper
+
+identity_wrapper = DynamoIdentityWrapper()
+
+
+LOG_COUNT_EVERY = 100_000
+
+
+logger = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    def handle(self, *args: Any, **options: Any) -> None:
+        total_count = identity_wrapper.table.item_count
+        scanned_count = 0
+        fixed_count = 0
+
+        for identity_document in identity_wrapper.query_get_all_items():
+            should_write_identity_document = False
+
+            if identity_traits_data := identity_document.get("identity_traits"):
+                blank_traits = (
+                    trait_data
+                    for trait_data in identity_traits_data
+                    if "trait_value" not in trait_data
+                )
+                for trait_data in blank_traits:
+                    should_write_identity_document = True
+                    trait_data["trait_value"] = ""
+
+            scanned_count += 1
+
+            if should_write_identity_document:
+                identity_wrapper.put_item(identity_document)
+                fixed_count += 1
+                self.stdout.write(
+                    "fixed identity "
+                    f"scanned={scanned_count}/{total_count} "
+                    f"percentage={scanned_count/total_count*100:.2f} "
+                    f"fixed={fixed_count} "
+                    f"id={identity_document['identity_uuid']}",
+                )
+
+            if not (scanned_count % LOG_COUNT_EVERY):
+                self.stdout.write(
+                    f"scanned={scanned_count}/{total_count} "
+                    f"percentage={scanned_count/total_count*100:.2f} "
+                    f"fixed={fixed_count}"
+                )
+
+        self.stdout.write(
+            self.style.SUCCESS(
+                "finished "
+                f"scanned={scanned_count}/{total_count} "
+                f"percentage={scanned_count/total_count*100:.2f} "
+                f"fixed={fixed_count}"
+            )
+        )

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/pyproject.toml

@@ -201,7 +201,7 @@ workflows-logic = { git = "https://github.com/flagsmith/flagsmith-workflows", ta
 [tool.poetry.group.dev.dependencies]
 django-test-migrations = "~1.2.0"
 responses = "~0.22.0"
-pre-commit = "~3.0.4"
+pre-commit = "^4.0.1"
 pytest-mock = "~3.10.0"
 pytest-lazy-fixture = "~0.6.3"
 moto = "~4.1.3"

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