feat(api): Add command to remove identity overrides that no longer exist

api/edge_api/identities/dataclasses.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from enum import Enum
+
+
+class OrphanedIdentityOverrideReason(Enum):
+    """
+    Why an `environments_v2` identity override no longer reflects its identity.
+    """
+
+    # No identity document exists for the identifier any more.
+    IDENTITY_DELETED = "identity_deleted"
+    # The identifier exists, but as a different identity than the one the
+    # override was written against, so the override can never be reached.
+    IDENTITY_UUID_CHANGED = "identity_uuid_changed"
+    # The identity is the same one, but no longer overrides this feature.
+    OVERRIDE_REMOVED = "override_removed"
+
+
+@dataclass(frozen=True)
+class OrphanedIdentityOverride:
+    document_key: str
+    identifier: str
+    identity_uuid: str
+    feature_id: int
+    feature_name: str
+    reason: OrphanedIdentityOverrideReason

api/edge_api/identities/edge_identity_service.py

@@ -1,10 +1,24 @@
+import typing
+from collections import defaultdict
+from typing import Any, Generator
+
 from django.core.exceptions import ObjectDoesNotExist
 
+from edge_api.identities.dataclasses import (
+    OrphanedIdentityOverride,
+    OrphanedIdentityOverrideReason,
+)
 from edge_api.identities.models import EdgeIdentity
 from environments.dynamodb import DynamoEnvironmentV2Wrapper
+from environments.dynamodb.constants import DYNAMODB_MAX_BATCH_GET_ITEM_COUNT
 from environments.dynamodb.types import (
     IdentityOverrideV2,
 )
+from util.engine_models.identities.models import IdentityModel
+from util.util import iter_chunks
+
+if typing.TYPE_CHECKING:
+    from environments.models import Environment
 
 ddb_environment_v2_wrapper = DynamoEnvironmentV2Wrapper()
 
@@ -41,6 +55,104 @@ def get_edge_identity_override_keys(environment_id: int) -> list[str]:
     return [item["document_key"] for item in override_items]
 
 
+def iter_orphaned_identity_overrides(
+    environment: "Environment",
+) -> Generator[OrphanedIdentityOverride, None, None]:
+    """
+    Yield the environment's identity overrides that its identities no longer have.
+
+    The identity document is the source of truth: it is what remote evaluation
+    and the identity page read. An `environments_v2` override the identity does
+    not have is therefore stale, and is still served to local evaluation SDKs
+    and listed on the feature's identity overrides tab.
+    """
+    override_documents = (
+        ddb_environment_v2_wrapper.get_identity_overrides_by_environment_id(
+            environment_id=environment.id,
+            projection_expression_attributes=[
+                "document_key",
+                "identifier",
+                "identity_uuid",
+                "feature_state.feature",
+            ],
+        )
+    )
+    documents_by_identifier: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for override_document in override_documents:
+        documents_by_identifier[override_document["identifier"]].append(
+            override_document
+        )
+
+    for identifiers in iter_chunks(
+        documents_by_identifier,
+        chunk_size=DYNAMODB_MAX_BATCH_GET_ITEM_COUNT,
+    ):
+        identifiers_by_composite_key = {
+            IdentityModel.generate_composite_key(
+                environment.api_key, identifier
+            ): identifier
+            for identifier in identifiers
+        }
+        identity_documents = {
+            identity_document["composite_key"]: identity_document
+            for identity_document in EdgeIdentity.dynamo_wrapper.iter_items_by_composite_keys(
+                identifiers_by_composite_key,
+                projection_expression="composite_key,identity_uuid,identity_features",
+            )
+        }
+        for composite_key, identifier in identifiers_by_composite_key.items():
+            identity_document = identity_documents.get(composite_key)
+            for override_document in documents_by_identifier[identifier]:
+                feature = override_document["feature_state"]["feature"]
+                if reason := _get_orphaned_identity_override_reason(
+                    override_document=override_document,
+                    identity_document=identity_document,
+                ):
+                    yield OrphanedIdentityOverride(
+                        document_key=override_document["document_key"],
+                        identifier=identifier,
+                        identity_uuid=override_document["identity_uuid"],
+                        feature_id=int(feature["id"]),
+                        feature_name=feature["name"],
+                        reason=reason,
+                    )
+
+
+def delete_orphaned_identity_override(
+    environment_id: int,
+    orphaned_identity_override: OrphanedIdentityOverride,
+) -> bool:
+    """
+    Delete a stale identity override, unless it has been rewritten since it was read.
+
+    :return: whether the override was deleted.
+    """
+    return ddb_environment_v2_wrapper.delete_identity_override_if_unchanged(
+        environment_id=environment_id,
+        document_key=orphaned_identity_override.document_key,
+        identity_uuid=orphaned_identity_override.identity_uuid,
+    )
+
+
+def _get_orphaned_identity_override_reason(
+    override_document: dict[str, Any],
+    identity_document: dict[str, Any] | None,
+) -> OrphanedIdentityOverrideReason | None:
+    if identity_document is None:
+        return OrphanedIdentityOverrideReason.IDENTITY_DELETED
+    if identity_document["identity_uuid"] != override_document["identity_uuid"]:
+        return OrphanedIdentityOverrideReason.IDENTITY_UUID_CHANGED
+    overridden_feature_ids = {
+        int(feature_state["feature"]["id"])
+        for feature_state in identity_document.get("identity_features") or []
+    }
+    if int(override_document["feature_state"]["feature"]["id"]) not in (
+        overridden_feature_ids
+    ):
+        return OrphanedIdentityOverrideReason.OVERRIDE_REMOVED
+    return None
+
+
 def get_overridden_feature_ids_for_edge_identity(identity_uuid: str) -> set[int]:
     try:
         identity_document = EdgeIdentity.dynamo_wrapper.get_item_from_uuid(

api/edge_api/management/commands/delete_orphaned_identity_overrides.py

@@ -0,0 +1,94 @@
+from argparse import ArgumentParser
+from collections import Counter
+from typing import Any
+
+import structlog
+from django.core.management import BaseCommand, CommandError
+
+from edge_api.identities.edge_identity_service import (
+    delete_orphaned_identity_override,
+    iter_orphaned_identity_overrides,
+)
+from environments.models import Environment
+
+logger: structlog.BoundLogger = structlog.get_logger("edge_identities")
+
+
+class Command(BaseCommand):
+    help = (
+        "Delete identity overrides in the environments_v2 table that the identity "
+        "they belong to no longer has, so that local evaluation and the dashboard "
+        "agree with the identity document."
+    )
+
+    def add_arguments(self, parser: ArgumentParser) -> None:
+        parser.add_argument(
+            "--environment-id",
+            dest="environment_id",
+            type=int,
+            required=True,
+            help="ID of the environment to reconcile",
+        )
+        parser.add_argument(
+            "--dry-run",
+            dest="dry_run",
+            action="store_true",
+            help="Report what would be deleted without deleting anything",
+        )
+
+    def handle(
+        self,
+        *args: Any,
+        environment_id: int,
+        dry_run: bool,
+        **options: Any,
+    ) -> None:
+        try:
+            environment = Environment.objects.get(id=environment_id)
+        except Environment.DoesNotExist:
+            raise CommandError(f"Environment {environment_id} does not exist")
+
+        log: structlog.BoundLogger = logger.bind(  # type: ignore[assignment]
+            environment__id=environment.id,
+            dry_run=dry_run,
+        )
+        log.info("identity_override.reconciliation_started")
+
+        reasons: Counter[str] = Counter()
+        deleted_count = skipped_count = 0
+
+        for orphaned_identity_override in iter_orphaned_identity_overrides(environment):
+            reasons[orphaned_identity_override.reason.value] += 1
+            self.stdout.write(
+                "\t".join(
+                    [
+                        orphaned_identity_override.reason.value,
+                        orphaned_identity_override.feature_name,
+                        orphaned_identity_override.identifier,
+                        orphaned_identity_override.document_key,
+                    ]
+                )
+            )
+            if dry_run:
+                continue
+            if delete_orphaned_identity_override(
+                environment_id=environment.id,
+                orphaned_identity_override=orphaned_identity_override,
+            ):
+                deleted_count += 1
+            else:
+                # The override was rewritten between being read and being
+                # deleted, so it is no longer the stale document we identified.
+                skipped_count += 1
+                log.info(
+                    "identity_override.delete_skipped",
+                    document_key=orphaned_identity_override.document_key,
+                )
+
+        log.info(
+            "identity_override.reconciliation_finished",
+            orphaned__count=sum(reasons.values()),
+            deleted__count=deleted_count,
+            skipped__count=skipped_count,
+            reasons=dict(reasons),
+        )

api/environments/dynamodb/constants.py

@@ -7,6 +7,7 @@
 ENVIRONMENTS_V2_SECONDARY_INDEX_PARTITION_KEY = "environment_api_key"
 
 DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT = 25
+DYNAMODB_MAX_BATCH_GET_ITEM_COUNT = 100
 IDENTITIES_PAGINATION_LIMIT = 1000
 
 SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS = 3

api/environments/dynamodb/wrappers/base.py

@@ -8,7 +8,7 @@
 from sentry_sdk import set_context  # TODO @kgustyr: Replace with OTel
 
 if typing.TYPE_CHECKING:
-    from mypy_boto3_dynamodb.service_resource import Table
+    from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table
     from mypy_boto3_dynamodb.type_defs import (
         QueryOutputTableTypeDef,
         ScanOutputTableTypeDef,
@@ -27,21 +27,33 @@ class BaseDynamoWrapper:
 
     def __init__(self) -> None:
         self._table: typing.Optional["Table"] = None
+        self._resource: typing.Optional["DynamoDBServiceResource"] = None
 
     @property
     def table(self) -> typing.Optional["Table"]:
         if not self._table:
             self._table = self.get_table()
         return self._table
 
+    @property
+    def resource(self) -> "DynamoDBServiceResource":
+        """
+        The service resource behind `table`, for operations that span items —
+        e.g. `batch_get_item`, which is not available on a `Table`.
+        """
+        if not self._resource:
+            self._resource = self.get_resource()
+        return self._resource
+
     def get_table_name(self) -> str:
         return self.table_name
 
+    def get_resource(self) -> "DynamoDBServiceResource":
+        return boto3.resource("dynamodb", config=Config(tcp_keepalive=True))
+
     def get_table(self) -> "Table | None":
         if table_name := self.get_table_name():
-            return boto3.resource("dynamodb", config=Config(tcp_keepalive=True)).Table(
-                table_name
-            )
+            return self.resource.Table(table_name)
         return None
 
     @property

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -3,7 +3,8 @@
 from typing import Any, Iterable
 
 import structlog
-from boto3.dynamodb.conditions import ConditionBase, Key
+from boto3.dynamodb.conditions import Attr, ConditionBase, Key
+from botocore.exceptions import ClientError
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
 from django.db.models import prefetch_related_objects
@@ -216,6 +217,35 @@ def delete_environment(self, environment_id: int):  # type: ignore[no-untyped-de
                     },
                 )
 
+    def delete_identity_override_if_unchanged(
+        self,
+        environment_id: int,
+        document_key: str,
+        identity_uuid: str,
+    ) -> bool:
+        """
+        Delete a single identity override document, unless it has been rewritten.
+
+        The condition guards against removing an override that was recreated
+        between a caller deciding it was stale and this delete landing — without
+        it, a repair would carry the same race as the bug it repairs.
+
+        :return: whether the document was deleted.
+        """
+        try:
+            self.table.delete_item(  # type: ignore[union-attr]
+                Key={
+                    ENVIRONMENTS_V2_PARTITION_KEY: str(environment_id),
+                    ENVIRONMENTS_V2_SORT_KEY: document_key,
+                },
+                ConditionExpression=Attr("identity_uuid").eq(identity_uuid),
+            )
+        except ClientError as exc:
+            if exc.response["Error"]["Code"] == "ConditionalCheckFailedException":
+                return False
+            raise
+        return True
+
     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