From 574a08e274b1ee8de56faa765b3bb9a4ec464473 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Tue, 12 Dec 2023 10:20:04 +0000 Subject: [PATCH] feat: Migrate given project's (edge) identities to environments v2 (#3138) Co-authored-by: Matthew Elwell --- api/conftest.py | 47 ++++++++++- api/environments/dynamodb/constants.py | 4 + api/environments/dynamodb/dynamodb_wrapper.py | 58 +++++++++++--- api/environments/dynamodb/services.py | 65 +++++++++++++++ api/projects/admin.py | 2 +- ...add_identity_overrides_migration_status.py | 8 +- api/projects/models.py | 26 ++++-- api/projects/tasks.py | 16 ++++ api/task_processor/decorators.py | 3 - ...st_unit_dynamodb_environment_v2_wrapper.py | 20 +++++ .../test_unit_dynamodb_identity_wrapper.py | 71 ++++++++++++++-- .../dynamodb/test_unit_services.py | 80 +++++++++++++++++++ api/tests/unit/projects/test_tasks.py | 70 ++++++++++++++++ .../projects/test_unit_projects_models.py | 13 +-- .../mappers/test_unit_mappers_dynamodb.py | 5 +- api/util/mappers/__init__.py | 2 + api/util/mappers/dynamodb.py | 5 +- 17 files changed, 453 insertions(+), 42 deletions(-) create mode 100644 api/environments/dynamodb/services.py create mode 100644 api/tests/unit/environments/dynamodb/test_unit_services.py create mode 100644 api/tests/unit/projects/test_tasks.py diff --git a/api/conftest.py b/api/conftest.py index 94c6668ce9e4..24651534dc31 100644 --- a/api/conftest.py +++ b/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, diff --git a/api/environments/dynamodb/constants.py b/api/environments/dynamodb/constants.py index 7f59ab6086e3..af4d6c39a4e7 100644 --- a/api/environments/dynamodb/constants.py +++ b/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 diff --git a/api/environments/dynamodb/dynamodb_wrapper.py b/api/environments/dynamodb/dynamodb_wrapper.py index 0158c775128d..324dffe9b64e 100644 --- a/api/environments/dynamodb/dynamodb_wrapper.py +++ b/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 diff --git a/api/environments/dynamodb/services.py b/api/environments/dynamodb/services.py new file mode 100644 index 000000000000..ffa3abb1b11d --- /dev/null +++ b/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), + ) diff --git a/api/projects/admin.py b/api/projects/admin.py index 08be898a629e..96e90d837370 100644 --- a/api/projects/admin.py +++ b/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( diff --git a/api/projects/migrations/0021_add_identity_overrides_migration_status.py b/api/projects/migrations/0021_add_identity_overrides_migration_status.py index 0e0c8cd136fc..e3f141fc53bb 100644 --- a/api/projects/migrations/0021_add_identity_overrides_migration_status.py +++ b/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"), diff --git a/api/projects/models.py b/api/projects/models.py index 27cc7d236c5b..639bab563eb9 100644 --- a/api/projects/models.py +++ b/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 ) diff --git a/api/projects/tasks.py b/api/projects/tasks.py index bda0ed7a3857..743d741c1016 100644 --- a/api/projects/tasks.py +++ b/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() diff --git a/api/task_processor/decorators.py b/api/task_processor/decorators.py index 6cfe235fe181..6819b85d1fd9 100644 --- a/api/task_processor/decorators.py +++ b/api/task_processor/decorators.py @@ -6,7 +6,6 @@ from threading import Thread from django.conf import settings -from django.core.serializers.json import DjangoJSONEncoder from django.db.transaction import on_commit from django.utils import timezone @@ -19,8 +18,6 @@ logger = logging.getLogger(__name__) -_django_json_encoder_default = DjangoJSONEncoder().default - class TaskHandler(typing.Generic[P]): __slots__ = ( diff --git a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py b/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py index 6e08775a1a32..52086b3a259b 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py +++ b/api/tests/unit/environments/dynamodb/test_unit_dynamodb_environment_v2_wrapper.py @@ -134,3 +134,23 @@ def test_environment_v2_wrapper__update_identity_overrides__delete_expected( # Then results = flagsmith_environments_v2_table.scan()["Items"] assert len(results) == 0 + + +def test_environment_v2_wrapper__write_environments__put_expected( + settings: SettingsWrapper, + environment: Environment, + flagsmith_environments_v2_table: Table, +) -> None: + # Given + settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name + wrapper = DynamoEnvironmentV2Wrapper() + + # When + wrapper.write_environments( + environments=[environment], + ) + + # Then + results = flagsmith_environments_v2_table.scan()["Items"] + assert len(results) == 1 + assert results[0] == map_environment_to_environment_v2_document(environment) diff --git a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py b/api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py index f1ca7411506b..169ccd29295f 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py +++ b/api/tests/unit/environments/dynamodb/test_unit_dynamodb_identity_wrapper.py @@ -4,7 +4,7 @@ from boto3.dynamodb.conditions import Key from core.constants import INTEGER from django.core.exceptions import ObjectDoesNotExist -from flag_engine.identities.builders import build_identity_model +from flag_engine.identities.models import IdentityModel from flag_engine.segments.constants import IN from rest_framework.exceptions import NotFound @@ -246,17 +246,14 @@ def test_is_enabled_is_false_if_dynamo_table_name_is_not_set(settings, mocker): def test_is_enabled_is_true_if_dynamo_table_name_is_set(settings, mocker): # Given table_name = "random_table_name" - mocker.patch( - "environments.dynamodb.dynamodb_wrapper.DynamoIdentityWrapper.table_name", - table_name, - ) + settings.IDENTITIES_TABLE_NAME_DYNAMO = table_name mocked_config = mocker.patch("environments.dynamodb.dynamodb_wrapper.Config") mocked_boto3 = mocker.patch("environments.dynamodb.dynamodb_wrapper.boto3") # When dynamo_identity_wrapper = DynamoIdentityWrapper() - # Then + # Then assert dynamo_identity_wrapper.is_enabled is True mocked_boto3.resource.assert_called_with( "dynamodb", config=mocked_config(tcp_keepalive=True) @@ -381,7 +378,7 @@ def test_get_segment_ids_with_identity_model(identity, environment, mocker): # Given dynamo_identity_wrapper = DynamoIdentityWrapper() identity_document = map_identity_to_identity_document(identity) - identity_model = build_identity_model(identity_document) + identity_model = IdentityModel.parse_obj(identity_document) mocker.patch.object( dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document @@ -398,3 +395,63 @@ def test_get_segment_ids_with_identity_model(identity, environment, mocker): # Then assert segment_ids == [] + + +def test_identity_wrapper__iter_all_items_paginated__returns_expected( + environment: "Environment", + identity: "Identity", + mocker: "MockerFixture", +) -> None: + # Given + dynamo_identity_wrapper = DynamoIdentityWrapper() + identity_document = map_identity_to_identity_document(identity) + environment_api_key = "test_api_key" + limit = 1 + + expected_engine_identity = IdentityModel.parse_obj(identity_document) + expected_next_page_key = "next_page_key" + + environment_document = map_environment_to_environment_document(environment) + mocked_environment_wrapper = mocker.patch( + "environments.dynamodb.dynamodb_wrapper.DynamoEnvironmentWrapper", + autospec=True, + ) + mocked_environment_wrapper.return_value.get_item.return_value = environment_document + + mocked_get_all_items = mocker.patch.object( + dynamo_identity_wrapper, + "get_all_items", + autospec=True, + ) + mocked_get_all_items.side_effect = [ + {"Items": [identity_document], "LastEvaluatedKey": "next_page_key"}, + {"Items": [identity_document], "LastEvaluatedKey": None}, + ] + + # When + iterator = dynamo_identity_wrapper.iter_all_items_paginated( + environment_api_key=environment_api_key, limit=limit + ) + result_1 = next(iterator) + result_2 = next(iterator) + + # Then + with pytest.raises(StopIteration): + next(iterator) + + assert result_1 == expected_engine_identity + assert result_2 == expected_engine_identity + + mocked_get_all_items.assert_has_calls( + [ + mocker.call( + environment_api_key=environment_api_key, + limit=limit, + ), + mocker.call( + environment_api_key=environment_api_key, + limit=limit, + start_key=expected_next_page_key, + ), + ] + ) diff --git a/api/tests/unit/environments/dynamodb/test_unit_services.py b/api/tests/unit/environments/dynamodb/test_unit_services.py new file mode 100644 index 000000000000..63fdd71bb152 --- /dev/null +++ b/api/tests/unit/environments/dynamodb/test_unit_services.py @@ -0,0 +1,80 @@ +from mypy_boto3_dynamodb.service_resource import Table +from pytest_mock import MockerFixture + +from environments.dynamodb.dynamodb_wrapper import ( + DynamoEnvironmentV2Wrapper, + DynamoIdentityWrapper, +) +from environments.dynamodb.services import migrate_environments_to_v2 +from environments.identities.models import Identity +from environments.models import Environment +from features.models import FeatureState +from util.mappers import ( + map_engine_feature_state_to_identity_override, + map_engine_identity_to_identity_document, + map_environment_to_environment_v2_document, + map_identity_override_to_identity_override_document, + map_identity_to_engine, +) + + +def test_migrate_environments_to_v2__environment_with_overrides__writes_expected( + environment: Environment, + identity: Identity, + identity_featurestate: FeatureState, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper, + flagsmith_environments_v2_table: Table, +) -> None: + # Given + engine_identity = map_identity_to_engine(identity, with_overrides=True) + dynamodb_identity_wrapper.put_item( + map_engine_identity_to_identity_document(engine_identity) + ) + + expected_environment_document = map_environment_to_environment_v2_document( + environment + ) + expected_identity_override_document = ( + map_identity_override_to_identity_override_document( + map_engine_feature_state_to_identity_override( + feature_state=engine_identity.identity_features[0], + identity_uuid=str(engine_identity.identity_uuid), + identifier=engine_identity.identifier, + environment_api_key=environment.api_key, + environment_id=environment.id, + ), + ) + ) + + # When + migrate_environments_to_v2(project_id=environment.project_id) + + # Then + results = flagsmith_environments_v2_table.scan()["Items"] + assert len(results) == 2 + assert results[0] == expected_environment_document + assert results[1] == expected_identity_override_document + + +def test_migrate_environments_to_v2__wrapper_disabled__does_not_write( + mocker: MockerFixture, +) -> None: + # Given + mocked_dynamodb_identity_wrapper = mocker.patch( + "environments.dynamodb.services.DynamoIdentityWrapper", + autospec=True, + return_value=mocker.MagicMock(is_enabled=False), + ) + mocked_dynamodb_v2_wrapper = mocker.patch( + "environments.dynamodb.services.DynamoEnvironmentV2Wrapper", + autospec=True, + return_value=mocker.MagicMock(is_enabled=False), + ) + + # When + migrate_environments_to_v2(project_id=mocker.Mock()) + + # Then + mocked_dynamodb_identity_wrapper.return_value.assert_not_called() + mocked_dynamodb_v2_wrapper.return_value.assert_not_called() diff --git a/api/tests/unit/projects/test_tasks.py b/api/tests/unit/projects/test_tasks.py new file mode 100644 index 000000000000..1812f5b54c2d --- /dev/null +++ b/api/tests/unit/projects/test_tasks.py @@ -0,0 +1,70 @@ +import pytest +from pytest_mock import MockerFixture + +from projects.models import IdentityOverridesV2MigrationStatus, Project +from projects.tasks import migrate_project_environments_to_v2 + + +@pytest.fixture +def project_v2_migration_in_progress( + project: Project, +) -> Project: + project.identity_overrides_v2_migration_status = ( + IdentityOverridesV2MigrationStatus.IN_PROGRESS + ) + project.save() + return project + + +def test_migrate_project_environments_to_v2__calls_expected( + mocker: MockerFixture, + project_v2_migration_in_progress: Project, +): + # Given + mocked_migrate_environments_to_v2 = mocker.patch( + "environments.dynamodb.services.migrate_environments_to_v2", + autospec=True, + return_value=None, + ) + + # When + migrate_project_environments_to_v2(project_id=project_v2_migration_in_progress.id) + + # Then + project_v2_migration_in_progress.refresh_from_db() + mocked_migrate_environments_to_v2.assert_called_once_with( + project_id=project_v2_migration_in_progress.id, + ) + assert project_v2_migration_in_progress.identity_overrides_v2_migration_status == ( + IdentityOverridesV2MigrationStatus.COMPLETE + ) + + +def test_migrate_project_environments_to_v2__expected_status_on_error( + mocker: MockerFixture, + project_v2_migration_in_progress: Project, +): + # Given + project_v2_migration_in_progress.identity_overrides_v2_migration_status = ( + IdentityOverridesV2MigrationStatus.IN_PROGRESS + ) + + mocked_migrate_environments_to_v2 = mocker.patch( + "environments.dynamodb.services.migrate_environments_to_v2", + autospec=True, + side_effect=Exception, + ) + + # When + with pytest.raises(Exception): + migrate_project_environments_to_v2( + project_id=project_v2_migration_in_progress.id + ) + + # Then + mocked_migrate_environments_to_v2.assert_called_once_with( + project_id=project_v2_migration_in_progress.id + ) + assert project_v2_migration_in_progress.identity_overrides_v2_migration_status == ( + IdentityOverridesV2MigrationStatus.IN_PROGRESS + ) diff --git a/api/tests/unit/projects/test_unit_projects_models.py b/api/tests/unit/projects/test_unit_projects_models.py index a69a4bc65bd6..38885c18a455 100644 --- a/api/tests/unit/projects/test_unit_projects_models.py +++ b/api/tests/unit/projects/test_unit_projects_models.py @@ -5,7 +5,7 @@ from django.conf import settings from django.utils import timezone -from projects.models import IdentityOverridesMigrationStatus, Project +from projects.models import IdentityOverridesV2MigrationStatus, Project now = timezone.now() tomorrow = now + timedelta(days=1) @@ -137,19 +137,20 @@ def test_environments_are_updated_in_dynamodb_when_project_id_updated( @pytest.mark.parametrize( - "identity_overrides_migration_status, expected_value", + "identity_overrides_v2_migration_status, expected_value", ( - (IdentityOverridesMigrationStatus.NOT_STARTED, False), - (IdentityOverridesMigrationStatus.COMPLETE, True), + (IdentityOverridesV2MigrationStatus.NOT_STARTED, False), + (IdentityOverridesV2MigrationStatus.IN_PROGRESS, False), + (IdentityOverridesV2MigrationStatus.COMPLETE, True), ), ) def test_show_edge_identity_overrides_for_feature( - identity_overrides_migration_status: IdentityOverridesMigrationStatus, + identity_overrides_v2_migration_status: IdentityOverridesV2MigrationStatus, expected_value: bool, ): assert ( Project( - identity_overrides_migration_status=identity_overrides_migration_status + identity_overrides_v2_migration_status=identity_overrides_v2_migration_status ).show_edge_identity_overrides_for_feature == expected_value ) diff --git a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py index cd1100b751f7..29011aab36a0 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py @@ -2,6 +2,9 @@ from decimal import Decimal from typing import TYPE_CHECKING +from environments.dynamodb.constants import ( + ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, +) from util.mappers import dynamodb if TYPE_CHECKING: # pragma: no cover @@ -143,7 +146,7 @@ def test_map_environment_to_environment_v2_document__call_expected( # Then assert result == { - "document_key": "META", + "document_key": ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, "environment_id": str(environment.id), "allow_client_traits": True, "amplitude_config": None, diff --git a/api/util/mappers/__init__.py b/api/util/mappers/__init__.py index 86695e498087..995f02308be2 100644 --- a/api/util/mappers/__init__.py +++ b/api/util/mappers/__init__.py @@ -11,6 +11,7 @@ from util.mappers.engine import ( map_feature_state_to_engine, map_feature_to_engine, + map_identity_to_engine, map_mv_option_to_engine, ) @@ -24,6 +25,7 @@ "map_feature_to_engine", "map_identity_changeset_to_identity_override_changeset", "map_identity_override_to_identity_override_document", + "map_identity_to_engine", "map_identity_to_identity_document", "map_mv_option_to_engine", ) diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index 5cae5dc0afae..e457192b2e71 100644 --- a/api/util/mappers/dynamodb.py +++ b/api/util/mappers/dynamodb.py @@ -6,6 +6,9 @@ from pydantic import BaseModel from edge_api.identities.types import IdentityChangeset +from environments.dynamodb.constants import ( + ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, +) from environments.dynamodb.types import ( IdentityOverridesV2Changeset, IdentityOverrideV2, @@ -54,7 +57,7 @@ def map_environment_to_environment_v2_document( ) -> Document: return { **map_environment_to_environment_document(environment), - "document_key": "META", + "document_key": ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, "environment_id": str(environment.id), }