feat: Add a task for writing (edge) identity overrides (#3127)

api/conftest.py

@@ -1,13 +1,19 @@
+import os
 import typing
 
+import boto3
 import pytest
 from django.contrib.contenttypes.models import ContentType
 from django.core.cache import cache
 from flag_engine.segments.constants import EQUAL
+from moto import mock_dynamodb
+from mypy_boto3_dynamodb.service_resource import 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
 from environments.identities.models import Identity
 from environments.identities.traits.models import Trait
 from environments.models import Environment, EnvironmentAPIKey
@@ -540,3 +546,51 @@ def project_content_type():
 @pytest.fixture
 def manage_user_group_permission(db):
     return OrganisationPermissionModel.objects.get(key=MANAGE_USER_GROUPS)
+
+
+@pytest.fixture()
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "eu-west-2"
+
+
+@pytest.fixture()
+def dynamodb(aws_credentials):
+    # TODO: move all wrapper tests to using moto
+    with mock_dynamodb():
+        yield boto3.resource("dynamodb")
+
+
+@pytest.fixture()
+def flagsmith_environments_v2_table(dynamodb) -> Table:
+    return dynamodb.create_table(
+        TableName="flagsmith_environments_v2",
+        KeySchema=[
+            {
+                "AttributeName": "environment_id",
+                "KeyType": "HASH",
+            },
+            {
+                "AttributeName": "document_key",
+                "KeyType": "RANGE",
+            },
+        ],
+        AttributeDefinitions=[
+            {"AttributeName": "environment_id", "AttributeType": "S"},
+            {"AttributeName": "document_key", "AttributeType": "S"},
+        ],
+        BillingMode="PAY_PER_REQUEST",
+    )
+
+
+@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/models.py

@@ -8,16 +8,20 @@
 from flag_engine.identities.models import IdentityFeaturesList, IdentityModel
 
 from api_keys.models import MasterAPIKey
+from edge_api.identities.tasks import (
+    generate_audit_log_records,
+    sync_identity_document_features,
+    update_flagsmith_environments_v2_identity_overrides,
+)
+from edge_api.identities.types import IdentityChangeset
+from edge_api.identities.utils import generate_change_dict
 from environments.dynamodb import DynamoIdentityWrapper
 from environments.models import Environment
 from features.models import FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
 from users.models import FFAdminUser
 from util.mappers import map_engine_identity_to_identity_document
 
-from .audit import generate_change_dict
-from .tasks import generate_audit_log_records, sync_identity_document_features
-
 
 class EdgeIdentity:
     dynamo_wrapper = DynamoIdentityWrapper()
@@ -161,7 +165,7 @@ def remove_feature_override(self, feature_state: FeatureStateModel) -> None:
 
     def save(self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None):
         self.dynamo_wrapper.put_item(self.to_document())
-        changes = self._get_changes(self._initial_state)
+        changes = self._get_changes()
         if changes["feature_overrides"]:
             # TODO: would this be simpler if we put a wrapper around FeatureStateModel instead?
             generate_audit_log_records.delay(
@@ -174,6 +178,13 @@ def save(self, user: FFAdminUser = None, master_api_key: MasterAPIKey = None):
                     "master_api_key_id": getattr(master_api_key, "id", None),
                 }
             )
+            update_flagsmith_environments_v2_identity_overrides.delay(
+                kwargs={
+                    "environment_api_key": self.environment_api_key,
+                    "changes": changes,
+                    "identity_uuid": str(self.identity_uuid),
+                }
+            )
         self._reset_initial_state()
 
     def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None:
@@ -187,7 +198,8 @@ def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> N
     def to_document(self) -> dict:
         return map_engine_identity_to_identity_document(self._engine_identity_model)
 
-    def _get_changes(self, previous_instance: "EdgeIdentity") -> dict:
+    def _get_changes(self) -> IdentityChangeset:
+        previous_instance = self._initial_state
         changes = {}
         feature_changes = changes.setdefault("feature_overrides", {})
         previous_feature_overrides = {
@@ -201,7 +213,9 @@ def _get_changes(self, previous_instance: "EdgeIdentity") -> dict:
             current_matching_fs = current_feature_overrides.get(uuid_)
             if current_matching_fs is None:
                 feature_changes[previous_fs.feature.name] = generate_change_dict(
-                    change_type="-", identity=self, old=previous_fs
+                    change_type="-",
+                    identity_id=self.id,
+                    old=previous_fs,
                 )
             elif (
                 current_matching_fs.enabled != previous_fs.enabled
@@ -210,15 +224,17 @@ def _get_changes(self, previous_instance: "EdgeIdentity") -> dict:
             ):
                 feature_changes[previous_fs.feature.name] = generate_change_dict(
                     change_type="~",
-                    identity=self,
+                    identity_id=self.id,
                     new=current_matching_fs,
                     old=previous_fs,
                 )
 
         for uuid_, previous_fs in current_feature_overrides.items():
             if uuid_ not in previous_feature_overrides:
                 feature_changes[previous_fs.feature.name] = generate_change_dict(
-                    change_type="+", identity=self, new=previous_fs
+                    change_type="+",
+                    identity_id=self.id,
+                    new=previous_fs,
                 )
 
         return changes

api/edge_api/identities/tasks.py

@@ -5,11 +5,14 @@
 
 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.models import Environment, Webhook
 from features.models import Feature, FeatureState
 from task_processor.decorators import register_task_handler
 from task_processor.models import TaskPriority
 from users.models import FFAdminUser
+from util.mappers import map_identity_changeset_to_identity_override_changeset
 from webhooks.webhooks import WebhookEventType, call_environment_webhooks
 
 logger = logging.getLogger(__name__)
@@ -89,21 +92,21 @@ def sync_identity_document_features(identity_uuid: str):
     )
 
     identity.synchronise_features(valid_feature_names)
-    EdgeIdentity.dynamo_wrapper.put_item(identity.to_document())
+    identity.save()
 
 
 @register_task_handler()
 def generate_audit_log_records(
     environment_api_key: str,
     identifier: str,
     identity_uuid: str,
-    changes: dict,
-    user_id: int = None,
-    master_api_key_id: int = None,
-):
+    changes: IdentityChangeset,
+    user_id: int | None = None,
+    master_api_key_id: int | None = None,
+) -> None:
     audit_records = []
 
-    feature_override_changes = changes.get("feature_overrides")
+    feature_override_changes = changes["feature_overrides"]
     if not feature_override_changes:
         return
 
@@ -113,7 +116,7 @@ def generate_audit_log_records(
 
     for feature_name, change_details in feature_override_changes.items():
         action = {"+": "created", "-": "deleted", "~": "updated"}.get(
-            change_details.get("change_type")
+            change_details["change_type"]
         )
         log = f"Feature override {action} for feature '{feature_name}' and identity '{identifier}'"
         audit_records.append(
@@ -130,3 +133,25 @@ def generate_audit_log_records(
         )
 
     AuditLog.objects.bulk_create(audit_records)
+
+
+@register_task_handler()
+def update_flagsmith_environments_v2_identity_overrides(
+    environment_api_key: str,
+    identity_uuid: str,
+    changes: IdentityChangeset,
+) -> None:
+    feature_override_changes = changes["feature_overrides"]
+    if not feature_override_changes:
+        return
+
+    environment = Environment.objects.get(api_key=environment_api_key)
+    dynamodb_wrapper_v2 = DynamoEnvironmentV2Wrapper()
+
+    identity_override_changeset = map_identity_changeset_to_identity_override_changeset(
+        identity_changeset=changes,
+        identity_uuid=identity_uuid,
+        environment_api_key=environment_api_key,
+        environment_id=environment.id,
+    )
+    dynamodb_wrapper_v2.update_identity_overrides(identity_override_changeset)

api/edge_api/identities/types.py

@@ -0,0 +1,15 @@
+from typing import Any, Literal, TypedDict
+
+from typing_extensions import NotRequired
+
+ChangeType = Literal["+", "-", "~"]
+
+
+class FeatureStateChangeDetails(TypedDict):
+    change_type: ChangeType
+    old: NotRequired[dict[str, Any]]
+    new: NotRequired[dict[str, Any]]
+
+
+class IdentityChangeset(TypedDict):
+    feature_overrides: dict[str, FeatureStateChangeDetails]

api/edge_api/identities/utils.py

@@ -0,0 +1,40 @@
+import typing
+
+from flag_engine.features.models import FeatureStateModel
+
+if typing.TYPE_CHECKING:
+    from edge_api.identities.types import ChangeType, FeatureStateChangeDetails
+
+
+def generate_change_dict(
+    change_type: "ChangeType",
+    identity_id: int | str | None,
+    new: FeatureStateModel | None = None,
+    old: FeatureStateModel | None = None,
+) -> "FeatureStateChangeDetails":
+    if not (new or old):
+        raise ValueError("Must provide one of 'new' or 'old'")
+
+    change_dict = {"change_type": change_type}
+    if new:
+        change_dict["new"] = _get_overridden_feature_state_dict(
+            identity_id=identity_id,
+            feature_state=new,
+        )
+    if old:
+        change_dict["old"] = _get_overridden_feature_state_dict(
+            identity_id=identity_id,
+            feature_state=old,
+        )
+
+    return change_dict
+
+
+def _get_overridden_feature_state_dict(
+    identity_id: int | str | None,
+    feature_state: FeatureStateModel,
+) -> dict[str, typing.Any]:
+    return {
+        **feature_state.dict(),
+        "feature_state_value": feature_state.get_value(identity_id),
+    }

api/environments/dynamodb/constants.py

@@ -0,0 +1,5 @@
+ENVIRONMENTS_V2_PARTITION_KEY = "environment_id"
+ENVIRONMENTS_V2_SORT_KEY = "document_key"
+
+ENVIRONMENTS_V2_SECONDARY_INDEX = "environment_api_key-index"
+ENVIRONMENTS_V2_SECONDARY_INDEX_PARTITION_KEY = "environment_api_key"