From 7e19f4ff2b2f3941e55360936b816879af4d06b6 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Mon, 25 Mar 2024 13:30:33 +0000 Subject: [PATCH] fix: No pagination when querying `environments_v2` (#3661) --- .../dynamodb/wrappers/environment_wrapper.py | 19 +++--- ...st_unit_dynamodb_environment_v2_wrapper.py | 67 +++++++++++++++++-- 2 files changed, 71 insertions(+), 15 deletions(-) diff --git a/api/environments/dynamodb/wrappers/environment_wrapper.py b/api/environments/dynamodb/wrappers/environment_wrapper.py index c8f40d2536d6..90e1bcd19fb9 100644 --- a/api/environments/dynamodb/wrappers/environment_wrapper.py +++ b/api/environments/dynamodb/wrappers/environment_wrapper.py @@ -67,17 +67,18 @@ def get_identity_overrides_by_environment_id( 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 list( + self.query_get_all_items( + 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 diff --git a/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py index 4e035f2d1f1b..c7324cc6c77b 100644 --- a/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py +++ b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py @@ -2,12 +2,16 @@ from mypy_boto3_dynamodb.service_resource import Table from pytest_django.fixtures import SettingsWrapper +from pytest_mock import MockerFixture from environments.dynamodb import DynamoEnvironmentV2Wrapper from environments.dynamodb.types import ( IdentityOverridesV2Changeset, IdentityOverrideV2, ) +from environments.dynamodb.utils import ( + get_environments_v2_identity_override_document_key, +) from environments.models import Environment from features.models import Feature, FeatureState from util.mappers import ( @@ -31,7 +35,9 @@ def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__retur identifier = "identity1" override_document = { "environment_id": str(environment.id), - "document_key": f"identity_override:{feature.id}:{identity_uuid}", + "document_key": get_environments_v2_identity_override_document_key( + feature_id=feature.id, identity_uuid=identity_uuid + ), "environment_api_key": environment.api_key, "identifier": identifier, "feature_state": {}, @@ -53,6 +59,45 @@ def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__retur assert results[0] == override_document +def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__last_evaluated_key__call_expected( + flagsmith_environments_v2_table: Table, + mocker: MockerFixture, +) -> None: + # Given + wrapper = DynamoEnvironmentV2Wrapper() + mocker.patch.object( + wrapper, "get_table" + ).return_value = table_mock = mocker.MagicMock(spec=flagsmith_environments_v2_table) + + last_evaluated_key = "next_page_key" + override_document = {"test": "document"} + + table_mock.query.side_effect = [ + {"Items": [override_document], "LastEvaluatedKey": last_evaluated_key}, + {"Items": [override_document], "LastEvaluatedKey": None}, + ] + + # When + results = wrapper.get_identity_overrides_by_environment_id( + environment_id=mocker.ANY, + feature_id=mocker.ANY, + ) + + # Then + assert results == [override_document, override_document] + wrapper.table.query.assert_has_calls( + [ + mocker.call( + KeyConditionExpression=mocker.ANY, + ), + mocker.call( + KeyConditionExpression=mocker.ANY, + ExclusiveStartKey=last_evaluated_key, + ), + ] + ) + + def test_environment_v2_wrapper__update_identity_overrides__put_expected( settings: SettingsWrapper, environment: Environment, @@ -69,7 +114,9 @@ def test_environment_v2_wrapper__update_identity_overrides__put_expected( override_document = IdentityOverrideV2.parse_obj( { "environment_id": str(environment.id), - "document_key": f"identity_override:{feature.id}:{identity_uuid}", + "document_key": get_environments_v2_identity_override_document_key( + feature_id=feature.id, identity_uuid=identity_uuid + ), "environment_api_key": environment.api_key, "feature_state": map_feature_state_to_engine(feature_state), "identifier": identifier, @@ -110,7 +157,9 @@ def test_environment_v2_wrapper__update_identity_overrides__delete_expected( IdentityOverrideV2.parse_obj( { "environment_id": str(environment.id), - "document_key": f"identity_override:{feature.id}:{identity_uuid}", + "document_key": get_environments_v2_identity_override_document_key( + feature_id=feature.id, identity_uuid=identity_uuid + ), "environment_api_key": environment.api_key, "feature_state": map_feature_state_to_engine(feature_state), "identifier": identifier, @@ -170,7 +219,9 @@ def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_d Item={ "environment_api_key": environment_api_key, "environment_id": environment_id, - "document_key": f"identity_override:{i}", + "document_key": get_environments_v2_identity_override_document_key( + feature_id=i, + ), } ) @@ -181,7 +232,9 @@ def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_d Item={ "environment_api_key": environment_2_api_key, "environment_id": environment_2_id, - "document_key": "identity_override:1", + "document_key": get_environments_v2_identity_override_document_key( + feature_id=1, + ), } ) @@ -194,5 +247,7 @@ def test_environment_v2_wrapper__delete_environment__deletes_related_data_from_d assert results[0] == { "environment_api_key": environment_2_api_key, "environment_id": environment_2_id, - "document_key": "identity_override:1", + "document_key": get_environments_v2_identity_override_document_key( + feature_id=1, + ), }