feat: Heal all identities with blank traits (#4908)

api/edge_api/management/commands/ensure_identity_traits_blanks.py

@@ -0,0 +1,65 @@
+from typing import Any
+
+from django.core.management import BaseCommand
+from structlog import get_logger
+from structlog.stdlib import BoundLogger
+
+from environments.dynamodb import DynamoIdentityWrapper
+
+identity_wrapper = DynamoIdentityWrapper()
+
+
+LOG_COUNT_EVERY = 100_000
+
+
+class Command(BaseCommand):
+    def handle(self, *args: Any, **options: Any) -> None:
+        total_count = identity_wrapper.table.item_count
+        scanned_count = 0
+        fixed_count = 0
+
+        log: BoundLogger = get_logger(total_count=total_count)
+        log.info("started")
+
+        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
+            scanned_percentage = scanned_count / total_count * 100
+
+            if should_write_identity_document:
+                identity_wrapper.put_item(identity_document)
+                fixed_count += 1
+
+                log.info(
+                    "identity-fixed",
+                    identity_uuid=identity_document["identity_uuid"],
+                    scanned_count=scanned_count,
+                    scanned_percentage=scanned_percentage,
+                    fixed_count=fixed_count,
+                )
+
+            if not (scanned_count % LOG_COUNT_EVERY):
+                log.info(
+                    "in-progress",
+                    scanned_count=scanned_count,
+                    scanned_percentage=scanned_percentage,
+                    fixed_count=fixed_count,
+                )
+
+        log.info(
+            "finished",
+            scanned_count=scanned_count,
+            scanned_percentage=scanned_percentage,
+            fixed_count=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

@@ -173,6 +173,7 @@ flagsmith-task-processor = { git = "https://github.com/Flagsmith/flagsmith-task-
 flagsmith-common = { git = "https://github.com/Flagsmith/flagsmith-common", tag = "v1.4.2" }
 tzdata = "^2024.1"
 djangorestframework-simplejwt = "^5.3.1"
+structlog = "^24.4.0"
 
 [tool.poetry.group.auth-controller]
 optional = true
@@ -201,7 +202,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"
@@ -219,6 +220,7 @@ requests-mock = "^1.11.0"
 django-extensions = "^3.2.3"
 pdbpp = "^0.10.3"
 mypy-boto3-dynamodb = "^1.33.0"
+pytest-structlog = "^1.1"
 
 [build-system]
 requires = ["poetry-core>=1.5.0"]