feat: Migrate given project's (edge) identities to environments v2 (#…

…3138)

Co-authored-by: Matthew Elwell <[email protected]>

api/conftest.py

@@ -7,13 +7,16 @@
 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 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
+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
@@ -566,7 +569,36 @@ def dynamodb(aws_credentials):
 
 
 @pytest.fixture()
-def flagsmith_environments_v2_table(dynamodb) -> Table:
+def flagsmith_identities_table(dynamodb: DynamoDBServiceResource) -> Table:
+    return dynamodb.create_table(
+        TableName="flagsmith_identities",
+        KeySchema=[
+            {
+                "AttributeName": "composite_key",
+                "KeyType": "HASH",
+            },
+        ],
+        AttributeDefinitions=[
+            {"AttributeName": "composite_key", "AttributeType": "S"},
+            {"AttributeName": "environment_api_key", "AttributeType": "S"},
+            {"AttributeName": "identifier", "AttributeType": "S"},
+        ],
+        GlobalSecondaryIndexes=[
+            {
+                "IndexName": "environment_api_key-identifier-index",
+                "KeySchema": [
+                    {"AttributeName": "environment_api_key", "KeyType": "HASH"},
+                    {"AttributeName": "identifier", "KeyType": "RANGE"},
+                ],
+                "Projection": {"ProjectionType": "ALL"},
+            }
+        ],
+        BillingMode="PAY_PER_REQUEST",
+    )
+
+
+@pytest.fixture()
+def flagsmith_environments_v2_table(dynamodb: DynamoDBServiceResource) -> Table:
     return dynamodb.create_table(
         TableName="flagsmith_environments_v2",
         KeySchema=[
@@ -587,6 +619,15 @@ def flagsmith_environments_v2_table(dynamodb) -> Table:
     )
 
 
+@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,

api/environments/dynamodb/constants.py

@@ -1,5 +1,9 @@
 ENVIRONMENTS_V2_PARTITION_KEY = "environment_id"
 ENVIRONMENTS_V2_SORT_KEY = "document_key"
 
+ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY = "_META"
+
 ENVIRONMENTS_V2_SECONDARY_INDEX = "environment_api_key-index"
 ENVIRONMENTS_V2_SECONDARY_INDEX_PARTITION_KEY = "environment_api_key"
+
+IDENTITIES_PAGINATION_LIMIT = 1000

api/environments/dynamodb/dynamodb_wrapper.py

@@ -14,9 +14,18 @@
 from flag_engine.segments.evaluator import get_identity_segments
 from rest_framework.exceptions import NotFound
 
+if typing.TYPE_CHECKING:
+    from mypy_boto3_dynamodb.service_resource import Table
+    from mypy_boto3_dynamodb.type_defs import (
+        QueryOutputTableTypeDef,
+        TableAttributeValueTypeDef,
+        QueryInputRequestTypeDef,
+    )
+
 from environments.dynamodb.constants import (
     ENVIRONMENTS_V2_PARTITION_KEY,
     ENVIRONMENTS_V2_SORT_KEY,
+    IDENTITIES_PAGINATION_LIMIT,
 )
 from environments.dynamodb.types import IdentityOverridesV2Changeset
 from environments.dynamodb.utils import (
@@ -25,6 +34,7 @@
 from util.mappers import (
     map_environment_api_key_to_environment_api_key_document,
     map_environment_to_environment_document,
+    map_environment_to_environment_v2_document,
     map_identity_override_to_identity_override_document,
     map_identity_to_identity_document,
 )
@@ -40,7 +50,7 @@ class BaseDynamoWrapper:
     table_name: str = None
 
     def __init__(self):
-        self._table = None
+        self._table: "Table" | None = None
         if table_name := self.get_table_name():
             self._table = boto3.resource(
                 "dynamodb", config=Config(tcp_keepalive=True)
@@ -55,9 +65,10 @@ def get_table_name(self):
 
 
 class DynamoIdentityWrapper(BaseDynamoWrapper):
-    table_name = settings.IDENTITIES_TABLE_NAME_DYNAMO
+    def get_table_name(self) -> str | None:
+        return settings.IDENTITIES_TABLE_NAME_DYNAMO
 
-    def query_items(self, *args, **kwargs):
+    def query_items(self, *args, **kwargs) -> "QueryOutputTableTypeDef":
         return self._table.query(*args, **kwargs)
 
     def put_item(self, identity_dict: dict):
@@ -101,18 +112,40 @@ def get_item_from_uuid_or_404(self, uuid: str) -> dict:
             raise NotFound() from e
 
     def get_all_items(
-        self, environment_api_key: str, limit: int, start_key: dict = None
-    ):
+        self,
+        environment_api_key: str,
+        limit: int,
+        start_key: dict[str, "TableAttributeValueTypeDef"] | None = None,
+    ) -> "QueryOutputTableTypeDef":
         filter_expression = Key("environment_api_key").eq(environment_api_key)
-        query_kwargs = {
+        query_kwargs: "QueryInputRequestTypeDef" = {
             "IndexName": "environment_api_key-identifier-index",
-            "Limit": limit,
             "KeyConditionExpression": filter_expression,
+            "Limit": limit,
         }
         if start_key:
-            query_kwargs.update(ExclusiveStartKey=start_key)
+            query_kwargs["ExclusiveStartKey"] = start_key
         return self.query_items(**query_kwargs)
 
+    def iter_all_items_paginated(
+        self,
+        environment_api_key: str,
+        limit: int = IDENTITIES_PAGINATION_LIMIT,
+    ) -> typing.Generator[IdentityModel, None, None]:
+        last_evaluated_key = "initial"
+        get_all_items_kwargs = {
+            "environment_api_key": environment_api_key,
+            "limit": limit,
+        }
+        while last_evaluated_key:
+            query_response = self.get_all_items(
+                **get_all_items_kwargs,
+            )
+            for item in query_response["Items"]:
+                yield IdentityModel.parse_obj(item)
+            if last_evaluated_key := query_response.get("LastEvaluatedKey"):
+                get_all_items_kwargs["start_key"] = last_evaluated_key
+
     def search_items_with_identifier(
         self,
         environment_api_key: str,
@@ -179,7 +212,7 @@ def get_item(self, api_key: str) -> dict:
 
 
 class DynamoEnvironmentV2Wrapper(BaseDynamoEnvironmentWrapper):
-    def get_table_name(self):
+    def get_table_name(self) -> str | None:
         return settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO
 
     def get_identity_overrides_by_feature_id(
@@ -221,6 +254,13 @@ def update_identity_overrides(
                     ),
                 )
 
+    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),
+                )
+
 
 class DynamoEnvironmentAPIKeyWrapper(BaseDynamoWrapper):
     table_name = settings.ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO

api/environments/dynamodb/services.py

@@ -0,0 +1,65 @@
+import logging
+from typing import Generator, Iterable
+
+from environments.dynamodb.dynamodb_wrapper import (
+    DynamoEnvironmentV2Wrapper,
+    DynamoIdentityWrapper,
+)
+from environments.dynamodb.types import (
+    IdentityOverridesV2Changeset,
+    IdentityOverrideV2,
+)
+from environments.models import Environment
+from util.mappers import map_engine_feature_state_to_identity_override
+
+logger = logging.getLogger(__name__)
+
+
+def migrate_environments_to_v2(project_id: int) -> None:
+    dynamo_wrapper_v2 = DynamoEnvironmentV2Wrapper()
+    identity_wrapper = DynamoIdentityWrapper()
+
+    if not (dynamo_wrapper_v2.is_enabled and identity_wrapper.is_enabled):
+        return
+
+    logger.info("Migrating environments to v2 for project %d", project_id)
+
+    environments_to_migrate = Environment.objects.filter(project_id=project_id)
+    identity_overrides_to_migrate = _iter_paginated_overrides(
+        identity_wrapper=identity_wrapper,
+        environments=environments_to_migrate,
+    )
+
+    changeset = IdentityOverridesV2Changeset(
+        to_put=list(identity_overrides_to_migrate), to_delete=[]
+    )
+    logger.info(
+        "Retrieved %d identity overrides to migrate for project %d",
+        len(changeset.to_put),
+        project_id,
+    )
+
+    dynamo_wrapper_v2.write_environments(environments_to_migrate)
+    dynamo_wrapper_v2.update_identity_overrides(changeset)
+
+    logger.info("Finished migrating environments to v2 for project %d", project_id)
+
+
+def _iter_paginated_overrides(
+    *,
+    identity_wrapper: DynamoIdentityWrapper,
+    environments: Iterable[Environment],
+) -> Generator[IdentityOverrideV2, None, None]:
+    for environment in environments:
+        environment_api_key = environment.api_key
+        for identity in identity_wrapper.iter_all_items_paginated(
+            environment_api_key=environment_api_key,
+        ):
+            for feature_state in identity.identity_features:
+                yield map_engine_feature_state_to_identity_override(
+                    feature_state=feature_state,
+                    identity_uuid=str(identity.identity_uuid),
+                    identifier=identity.identifier,
+                    environment_api_key=environment_api_key,
+                    environment_id=str(environment.id),
+                )

api/projects/admin.py

@@ -69,7 +69,7 @@ class ProjectAdmin(admin.ModelAdmin):
         "max_segments_allowed",
         "max_features_allowed",
         "max_segment_overrides_allowed",
-        "identity_overrides_migration_status",
+        "identity_overrides_v2_migration_status",
     )
 
     @admin.action(

api/projects/migrations/0021_add_identity_overrides_migration_status.py

@@ -3,12 +3,12 @@
 from django.db import migrations, models
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 
-from projects.models import IdentityOverridesMigrationStatus
+from projects.models import IdentityOverridesV2MigrationStatus
 
 
 def apply_defaults(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
     apps.get_model("projects", "Project").objects.all().update(
-        identity_overrides_migration_status=IdentityOverridesMigrationStatus.NOT_STARTED
+        identity_overrides_v2_migration_status=IdentityOverridesV2MigrationStatus.NOT_STARTED
     )
 
 
@@ -20,7 +20,7 @@ class Migration(migrations.Migration):
     operations = [
         migrations.AddField(
             model_name="project",
-            name="identity_overrides_migration_status",
+            name="identity_overrides_v2_migration_status",
             field=models.CharField(
                 choices=[
                     ("NOT_STARTED", "Not Started"),
@@ -35,7 +35,7 @@ class Migration(migrations.Migration):
         migrations.RunPython(apply_defaults, reverse_code=migrations.RunPython.noop),
         migrations.AlterField(
             model_name="project",
-            name="identity_overrides_migration_status",
+            name="identity_overrides_v2_migration_status",
             field=models.CharField(
                 choices=[
                     ("NOT_STARTED", "Not Started"),

api/projects/models.py

@@ -24,13 +24,16 @@
     PermissionModel,
 )
 from projects.managers import ProjectManager
-from projects.tasks import write_environments_to_dynamodb
+from projects.tasks import (
+    migrate_project_environments_to_v2,
+    write_environments_to_dynamodb,
+)
 
 project_segments_cache = caches[settings.PROJECT_SEGMENTS_CACHE_LOCATION]
 environment_cache = caches[settings.ENVIRONMENT_CACHE_NAME]
 
 
-class IdentityOverridesMigrationStatus(models.TextChoices):
+class IdentityOverridesV2MigrationStatus(models.TextChoices):
     NOT_STARTED = "NOT_STARTED", "Not Started"
     IN_PROGRESS = "IN_PROGRESS", "In Progress"
     COMPLETE = "COMPLETE", "Complete"
@@ -77,10 +80,10 @@ class Project(LifecycleModelMixin, SoftDeleteExportableModel):
         default=100,
         help_text="Max segments overrides allowed for any (one) environment within this project",
     )
-    identity_overrides_migration_status = models.CharField(
+    identity_overrides_v2_migration_status = models.CharField(
         max_length=50,
-        choices=IdentityOverridesMigrationStatus.choices,
-        default=IdentityOverridesMigrationStatus.NOT_STARTED,
+        choices=IdentityOverridesV2MigrationStatus.choices,
+        default=IdentityOverridesV2MigrationStatus.NOT_STARTED,
     )
 
     objects = ProjectManager()
@@ -136,6 +139,15 @@ def clear_environments_cache(self):
             list(self.environments.values_list("api_key", flat=True))
         )
 
+    @hook(
+        AFTER_SAVE,
+        when="identity_overrides_v2_migration_status",
+        has_changed=True,
+        is_now=IdentityOverridesV2MigrationStatus.IN_PROGRESS,
+    )
+    def trigger_environments_v2_migration(self) -> None:
+        migrate_project_environments_to_v2.delay(kwargs={"project_id": self.id})
+
     @hook(AFTER_UPDATE)
     def write_to_dynamo(self):
         write_environments_to_dynamodb.delay(kwargs={"project_id": self.id})
@@ -163,8 +175,8 @@ def is_feature_name_valid(self, feature_name: str) -> bool:
     @property
     def show_edge_identity_overrides_for_feature(self) -> bool:
         return (
-            self.identity_overrides_migration_status
-            == IdentityOverridesMigrationStatus.COMPLETE
+            self.identity_overrides_v2_migration_status
+            == IdentityOverridesV2MigrationStatus.COMPLETE
         )
 
 

api/projects/tasks.py

@@ -1,3 +1,5 @@
+from django.db import transaction
+
 from task_processor.decorators import register_task_handler
 
 
@@ -6,3 +8,17 @@ def write_environments_to_dynamodb(project_id: int):
     from environments.models import Environment
 
     Environment.write_environments_to_dynamodb(project_id=project_id)
+
+
+@register_task_handler()
+def migrate_project_environments_to_v2(project_id: int):
+    from environments.dynamodb.services import migrate_environments_to_v2
+    from projects.models import IdentityOverridesV2MigrationStatus, Project
+
+    with transaction.atomic():
+        project = Project.objects.select_for_update().get(id=project_id)
+        migrate_environments_to_v2(project_id=project_id)
+        project.identity_overrides_v2_migration_status = (
+            IdentityOverridesV2MigrationStatus.COMPLETE
+        )
+        project.save()