Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: No pagination when querying environments_v2 #3661

Merged
merged 2 commits into from
Mar 25, 2024
Merged

fix: No pagination when querying environments_v2 #3661

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
053abbf
fix: No pagination when querying `environments_v2`
khvn26 Mar 25, 2024
9ffb9d5
add tests, use util for key generation
khvn26 Mar 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions api/environments/dynamodb/wrappers/environment_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
67 changes: 61 additions & 6 deletions api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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": {},
Expand All @@ -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,
),
]
)
Comment on lines +87 to +98
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So just to make sure I'm understanding this correctly: All the pages are returned in a single call. The response to the client could potentially span multiple documents that could be quite long?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep. We have a customer with thousands of identity overrides for a single feature which prompted us to come up with this fix



def test_environment_v2_wrapper__update_identity_overrides__put_expected(
settings: SettingsWrapper,
environment: Environment,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
),
}
)

Expand All @@ -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,
),
}
)

Expand All @@ -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,
),
}