feat(dynamo_documents): propagate delete to dynamo (#3220)

api/conftest.py

@@ -8,15 +8,10 @@
 from flag_engine.segments.constants import EQUAL
 from moto import mock_dynamodb
 from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table
-from pytest_django.fixtures import SettingsWrapper
 from rest_framework.authtoken.models import Token
 from rest_framework.test import APIClient
 
 from api_keys.models import MasterAPIKey
-from environments.dynamodb.dynamodb_wrapper import (
-    DynamoEnvironmentV2Wrapper,
-    DynamoIdentityWrapper,
-)
 from environments.identities.models import Identity
 from environments.identities.traits.models import Trait
 from environments.models import Environment, EnvironmentAPIKey
@@ -618,21 +613,3 @@ def flagsmith_environments_v2_table(dynamodb: DynamoDBServiceResource) -> Table:
         ],
         BillingMode="PAY_PER_REQUEST",
     )
-
-
-@pytest.fixture
-def dynamodb_identity_wrapper(
-    settings: SettingsWrapper,
-    flagsmith_identities_table: Table,
-) -> DynamoIdentityWrapper:
-    settings.IDENTITIES_TABLE_NAME_DYNAMO = flagsmith_identities_table.name
-    return DynamoIdentityWrapper()
-
-
-@pytest.fixture
-def dynamodb_wrapper_v2(
-    settings: SettingsWrapper,
-    flagsmith_environments_v2_table: Table,
-) -> DynamoEnvironmentV2Wrapper:
-    settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name
-    return DynamoEnvironmentV2Wrapper()

api/edge_api/identities/edge_identity_service.py

@@ -1,6 +1,6 @@
 import typing
 
-from environments.dynamodb.dynamodb_wrapper import DynamoEnvironmentV2Wrapper
+from environments.dynamodb import DynamoEnvironmentV2Wrapper
 from environments.dynamodb.types import IdentityOverrideV2
 
 ddb_environment_v2_wrapper = DynamoEnvironmentV2Wrapper()

api/edge_api/identities/tasks.py

@@ -6,7 +6,7 @@
 from audit.models import AuditLog
 from audit.related_object_type import RelatedObjectType
 from edge_api.identities.types import IdentityChangeset
-from environments.dynamodb.dynamodb_wrapper import DynamoEnvironmentV2Wrapper
+from environments.dynamodb import DynamoEnvironmentV2Wrapper
 from environments.models import Environment, Webhook
 from features.models import Feature, FeatureState
 from task_processor.decorators import register_task_handler

api/environments/dynamodb/__init__.py

@@ -1,4 +1,5 @@
-from .dynamodb_wrapper import (
+from .types import DynamoProjectMetadata
+from .wrappers import (
     DynamoEnvironmentAPIKeyWrapper,
     DynamoEnvironmentV2Wrapper,
     DynamoEnvironmentWrapper,
@@ -10,4 +11,5 @@
     "DynamoEnvironmentV2Wrapper",
     "DynamoEnvironmentWrapper",
     "DynamoIdentityWrapper",
+    "DynamoProjectMetadata",
 )

api/environments/dynamodb/migrator.py

@@ -8,12 +8,12 @@
 from projects.models import Project
 from util.queryset import iterator_with_prefetch
 
-from .dynamodb_wrapper import (
+from .types import DynamoProjectMetadata, ProjectIdentityMigrationStatus
+from .wrappers import (
     DynamoEnvironmentAPIKeyWrapper,
     DynamoEnvironmentWrapper,
     DynamoIdentityWrapper,
 )
-from .types import DynamoProjectMetadata, ProjectIdentityMigrationStatus
 
 
 class IdentityMigrator:

api/environments/dynamodb/services.py

@@ -1,7 +1,9 @@
 import logging
 from typing import Generator, Iterable
 
-from environments.dynamodb.dynamodb_wrapper import (
+from flag_engine.identities.models import IdentityModel
+
+from environments.dynamodb import (
     DynamoEnvironmentV2Wrapper,
     DynamoIdentityWrapper,
 )
@@ -53,9 +55,10 @@ def _iter_paginated_overrides(
 ) -> Generator[IdentityOverrideV2, None, None]:
     for environment in environments:
         environment_api_key = environment.api_key
-        for identity in identity_wrapper.iter_all_items_paginated(
+        for item in identity_wrapper.iter_all_items_paginated(
             environment_api_key=environment_api_key,
         ):
+            identity = IdentityModel.parse_obj(item)
             for feature_state in identity.identity_features:
                 yield map_engine_feature_state_to_identity_override(
                     feature_state=feature_state,

api/environments/dynamodb/types.py

@@ -72,6 +72,10 @@ def finish_identity_migration(self):
     def _save(self):
         return project_metadata_table.put_item(Item=asdict(self))
 
+    def delete(self):
+        if project_metadata_table:
+            project_metadata_table.delete_item(Key={"id": self.id})
+
 
 class IdentityOverrideV2(BaseModel):
     environment_id: str

api/environments/dynamodb/wrappers/__init__.py

@@ -0,0 +1,13 @@
+from .environment_api_key_wrapper import DynamoEnvironmentAPIKeyWrapper
+from .environment_wrapper import (
+    DynamoEnvironmentV2Wrapper,
+    DynamoEnvironmentWrapper,
+)
+from .identity_wrapper import DynamoIdentityWrapper
+
+__all__ = (
+    "DynamoEnvironmentAPIKeyWrapper",
+    "DynamoEnvironmentV2Wrapper",
+    "DynamoEnvironmentWrapper",
+    "DynamoIdentityWrapper",
+)

api/environments/dynamodb/wrappers/base.py

@@ -0,0 +1,45 @@
+import typing
+
+import boto3
+from botocore.config import Config
+
+if typing.TYPE_CHECKING:
+    from mypy_boto3_dynamodb.service_resource import Table
+
+
+class BaseDynamoWrapper:
+    table_name: str = None
+
+    def __init__(self) -> None:
+        self._table: typing.Optional["Table"] = None
+
+    @property
+    def table(self) -> typing.Optional["Table"]:
+        if not self._table:
+            self._table = self.get_table()
+        return self._table
+
+    def get_table_name(self) -> str:
+        return self.table_name
+
+    def get_table(self) -> typing.Optional["Table"]:
+        if table_name := self.get_table_name():
+            return boto3.resource("dynamodb", config=Config(tcp_keepalive=True)).Table(
+                table_name
+            )
+
+    @property
+    def is_enabled(self) -> bool:
+        return self.table is not None
+
+    def query_get_all_items(self, **kwargs: dict) -> typing.Generator[dict, None, None]:
+        while True:
+            query_response = self.table.query(**kwargs)
+            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

api/environments/dynamodb/wrappers/environment_api_key_wrapper.py

@@ -0,0 +1,31 @@
+import typing
+
+from django.conf import settings
+
+from util.mappers import (
+    map_environment_api_key_to_environment_api_key_document,
+)
+
+from .base import BaseDynamoWrapper
+
+if typing.TYPE_CHECKING:
+    from environments.models import EnvironmentAPIKey
+
+
+class DynamoEnvironmentAPIKeyWrapper(BaseDynamoWrapper):
+    table_name = settings.ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO
+
+    def write_api_key(self, api_key: "EnvironmentAPIKey"):
+        self.write_api_keys([api_key])
+
+    def write_api_keys(self, api_keys: typing.Iterable["EnvironmentAPIKey"]):
+        with self.table.batch_writer() as writer:
+            for api_key in api_keys:
+                writer.put_item(
+                    Item=map_environment_api_key_to_environment_api_key_document(
+                        api_key
+                    )
+                )
+
+    def delete_api_key(self, key: str) -> None:
+        self.table.delete_item(Key={"key": key})

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -0,0 +1,129 @@
+import typing
+from typing import Any, Iterable
+
+from boto3.dynamodb.conditions import Key
+from django.conf import settings
+from django.core.exceptions import ObjectDoesNotExist
+
+from environments.dynamodb.constants import (
+    DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT,
+    ENVIRONMENTS_V2_PARTITION_KEY,
+    ENVIRONMENTS_V2_SORT_KEY,
+)
+from environments.dynamodb.types import IdentityOverridesV2Changeset
+from environments.dynamodb.utils import (
+    get_environments_v2_identity_override_document_key,
+)
+from util.mappers import (
+    map_environment_to_environment_document,
+    map_environment_to_environment_v2_document,
+    map_identity_override_to_identity_override_document,
+)
+from util.util import iter_paired_chunks
+
+from .base import BaseDynamoWrapper
+
+if typing.TYPE_CHECKING:
+    from mypy_boto3_dynamodb.type_defs import QueryInputRequestTypeDef
+
+    from environments.models import Environment
+
+
+class BaseDynamoEnvironmentWrapper(BaseDynamoWrapper):
+    def write_environment(self, environment: "Environment") -> None:
+        self.write_environments([environment])
+
+    def write_environments(self, environments: Iterable["Environment"]) -> None:
+        raise NotImplementedError()
+
+
+class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper):
+    table_name = settings.ENVIRONMENTS_TABLE_NAME_DYNAMO
+
+    def write_environments(self, environments: Iterable["Environment"]):
+        with self.table.batch_writer() as writer:
+            for environment in environments:
+                writer.put_item(
+                    Item=map_environment_to_environment_document(environment),
+                )
+
+    def get_item(self, api_key: str) -> dict:
+        try:
+            return self.table.get_item(Key={"api_key": api_key})["Item"]
+        except KeyError as e:
+            raise ObjectDoesNotExist() from e
+
+    def delete_environment(self, api_key: str) -> None:
+        self.table.delete_item(Key={"api_key": api_key})
+
+
+class DynamoEnvironmentV2Wrapper(BaseDynamoEnvironmentWrapper):
+    def get_table_name(self) -> str | None:
+        return settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO
+
+    def get_identity_overrides_by_environment_id(
+        self,
+        environment_id: int,
+        feature_id: int | None = None,
+    ) -> typing.List[dict[str, Any]]:
+        try:
+            response = self.table.query(
+                KeyConditionExpression=Key(ENVIRONMENTS_V2_PARTITION_KEY).eq(
+                    str(environment_id),
+                )
+                & Key(ENVIRONMENTS_V2_SORT_KEY).begins_with(
+                    get_environments_v2_identity_override_document_key(
+                        feature_id=feature_id,
+                    ),
+                )
+            )
+            return response["Items"]
+        except KeyError as e:
+            raise ObjectDoesNotExist() from e
+
+    def update_identity_overrides(
+        self,
+        changeset: IdentityOverridesV2Changeset,
+    ) -> None:
+        for to_put, to_delete in iter_paired_chunks(
+            changeset.to_put,
+            changeset.to_delete,
+            chunk_size=DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT,
+        ):
+            with self.table.batch_writer() as writer:
+                for identity_override_to_delete in to_delete:
+                    writer.delete_item(
+                        Key={
+                            ENVIRONMENTS_V2_PARTITION_KEY: identity_override_to_delete.environment_id,
+                            ENVIRONMENTS_V2_SORT_KEY: identity_override_to_delete.document_key,
+                        },
+                    )
+                for identity_override_to_put in to_put:
+                    writer.put_item(
+                        Item=map_identity_override_to_identity_override_document(
+                            identity_override_to_put
+                        ),
+                    )
+
+    def write_environments(self, environments: Iterable["Environment"]) -> None:
+        with self.table.batch_writer() as writer:
+            for environment in environments:
+                writer.put_item(
+                    Item=map_environment_to_environment_v2_document(environment),
+                )
+
+    def delete_environment(self, environment_id: int):
+        environment_id = str(environment_id)
+        filter_expression = Key(ENVIRONMENTS_V2_PARTITION_KEY).eq(environment_id)
+        query_kwargs: "QueryInputRequestTypeDef" = {
+            "KeyConditionExpression": filter_expression,
+            "ProjectionExpression": "document_key",
+        }
+        with self.table.batch_writer() as writer:
+            for item in self.query_get_all_items(**query_kwargs):
+                writer.delete_item(
+                    Key={
+                        ENVIRONMENTS_V2_PARTITION_KEY: environment_id,
+                        ENVIRONMENTS_V2_SORT_KEY: item["document_key"],
+                    },
+                )