diff --git a/api/edge_api/identities/edge_identity_service.py b/api/edge_api/identities/edge_identity_service.py index 3c1edc453252..7e18224731ab 100644 --- a/api/edge_api/identities/edge_identity_service.py +++ b/api/edge_api/identities/edge_identity_service.py @@ -1,7 +1,8 @@ -import typing - from environments.dynamodb import DynamoEnvironmentV2Wrapper -from environments.dynamodb.types import IdentityOverrideV2 +from environments.dynamodb.types import ( + IdentityOverridesV2List, + IdentityOverrideV2, +) ddb_environment_v2_wrapper = DynamoEnvironmentV2Wrapper() @@ -9,10 +10,11 @@ def get_edge_identity_overrides( environment_id: int, feature_id: int | None = None, -) -> typing.List[IdentityOverrideV2]: +) -> list[IdentityOverrideV2]: override_items = ( ddb_environment_v2_wrapper.get_identity_overrides_by_environment_id( - environment_id=environment_id, feature_id=feature_id + environment_id=environment_id, + feature_id=feature_id, ) ) return [ @@ -21,3 +23,33 @@ def get_edge_identity_overrides( ) for item in override_items ] + + +def get_edge_identity_overrides_for_feature_ids( + environment_id: int, + feature_ids: None | list[int] = None, +) -> list[IdentityOverridesV2List]: + query_responses = ( + ddb_environment_v2_wrapper.get_identity_overrides_by_environment_id( + environment_id=environment_id, + feature_ids=feature_ids, + ) + ) + + results = [] + for identity_overrides_query_response in query_responses: + identity_overrides = [ + IdentityOverrideV2.model_validate( + {**item, "environment_id": str(item["environment_id"])} + ) + for item in identity_overrides_query_response.items + ] + complete = identity_overrides_query_response.is_num_identity_overrides_complete + results.append( + IdentityOverridesV2List( + identity_overrides=identity_overrides, + is_num_identity_overrides_complete=complete, + ) + ) + + return results diff --git a/api/environments/dynamodb/types.py b/api/environments/dynamodb/types.py index 28edc8c31556..604e195b603b 100644 --- a/api/environments/dynamodb/types.py +++ b/api/environments/dynamodb/types.py @@ -96,6 +96,12 @@ class IdentityOverridesV2Changeset: to_put: list[IdentityOverrideV2] +@dataclass +class IdentityOverridesV2List: + identity_overrides: list[IdentityOverrideV2] + is_num_identity_overrides_complete: bool + + @dataclass class EdgeV2MigrationResult: identity_overrides_changeset: IdentityOverridesV2Changeset diff --git a/api/environments/dynamodb/wrappers/environment_wrapper.py b/api/environments/dynamodb/wrappers/environment_wrapper.py index ff3579b17d67..6ccb3e911798 100644 --- a/api/environments/dynamodb/wrappers/environment_wrapper.py +++ b/api/environments/dynamodb/wrappers/environment_wrapper.py @@ -1,4 +1,6 @@ import typing +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from typing import Any, Iterable from boto3.dynamodb.conditions import Key @@ -29,6 +31,12 @@ from environments.models import Environment +@dataclass +class IdentityOverridesQueryResponse: + items: list[dict[str, Any]] + is_num_identity_overrides_complete: bool + + class BaseDynamoEnvironmentWrapper(BaseDynamoWrapper): def write_environment(self, environment: "Environment") -> None: self.write_environments([environment]) @@ -66,23 +74,66 @@ def get_identity_overrides_by_environment_id( self, environment_id: int, feature_id: int | None = None, - ) -> typing.List[dict[str, Any]]: + feature_ids: None | list[int] = None, + ) -> list[dict[str, Any]] | list[IdentityOverridesQueryResponse]: + try: - return list( - self.query_iter_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( + if feature_ids is None: + return list( + self.query_iter_all_items( + KeyConditionExpression=self.get_identity_overrides_key_condition_expression( + environment_id=environment_id, feature_id=feature_id, - ), + ) ) ) - ) + + else: + futures = [] + with ThreadPoolExecutor() as executor: + for feature_id in feature_ids: + futures.append( + executor.submit( + self.get_identity_overrides_page, + environment_id, + feature_id, + ) + ) + + results = [future.result() for future in futures] + return results + except KeyError as e: raise ObjectDoesNotExist() from e + def get_identity_overrides_page( + self, environment_id: int, feature_id: int + ) -> IdentityOverridesQueryResponse: + query_response = self.table.query( + KeyConditionExpression=self.get_identity_overrides_key_condition_expression( + environment_id=environment_id, + feature_id=feature_id, + ) + ) + last_evaluated_key = query_response.get("LastEvaluatedKey") + return IdentityOverridesQueryResponse( + items=query_response["Items"], + is_num_identity_overrides_complete=last_evaluated_key is None, + ) + + def get_identity_overrides_key_condition_expression( + self, + environment_id: int, + feature_id: None | int, + ) -> Key: + return 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, + ), + ) + def update_identity_overrides( self, changeset: IdentityOverridesV2Changeset, diff --git a/api/features/dataclasses.py b/api/features/dataclasses.py index 45dc1f034fd0..6ff4306be638 100644 --- a/api/features/dataclasses.py +++ b/api/features/dataclasses.py @@ -14,6 +14,7 @@ class EnvironmentFeatureOverridesData: num_segment_overrides: int = 0 num_identity_overrides: typing.Optional[int] = None + is_num_identity_overrides_complete: bool = True def add_identity_override(self): """ diff --git a/api/features/features_service.py b/api/features/features_service.py index 3150f2badafc..8cc5b33d95ed 100644 --- a/api/features/features_service.py +++ b/api/features/features_service.py @@ -2,7 +2,7 @@ from concurrent.futures import ThreadPoolExecutor from edge_api.identities.edge_identity_service import ( - get_edge_identity_overrides, + get_edge_identity_overrides_for_feature_ids, ) from features.dataclasses import EnvironmentFeatureOverridesData from features.versioning.versioning_service import get_environment_flags_list @@ -16,12 +16,13 @@ def get_overrides_data( environment: "Environment", + feature_ids: None | list[int] = None, ) -> OverridesData: """ Get correct overrides counts for a given environment. :param project: project to get overrides data for - :return: overrides data getter + :return: overrides data getter dictionary of {feature_id: EnvironmentFeatureOverridesData} """ project = environment.project @@ -29,7 +30,7 @@ def get_overrides_data( if project.edge_v2_identity_overrides_migrated: # If v2 migration is complete, count segment overrides from Core # and identity overrides from DynamoDB. - return get_edge_overrides_data(environment) + return get_edge_overrides_data(environment, feature_ids) # If v2 migration is not started, in progress, or incomplete, # only count segment overrides from Core. # v1 Edge identity overrides are uncountable. @@ -71,7 +72,7 @@ def get_core_overrides_data( def get_edge_overrides_data( - environment: "Environment", + environment: "Environment", feature_ids: None | list[int] = None ) -> OverridesData: """ Get the number of identity / segment overrides in a given environment for each feature in the @@ -81,14 +82,18 @@ def get_edge_overrides_data( :param environment: the environment to get the overrides data for :return OverridesData: dictionary of {feature_id: EnvironmentFeatureOverridesData} """ + + assert feature_ids is not None + with ThreadPoolExecutor() as executor: get_environment_flags_list_future = executor.submit( get_environment_flags_list, environment, ) get_overrides_data_future = executor.submit( - get_edge_identity_overrides, + get_edge_identity_overrides_for_feature_ids, environment_id=environment.id, + feature_ids=feature_ids, ) all_overrides_data: OverridesData = {} @@ -98,11 +103,18 @@ def get_edge_overrides_data( ) if feature_state.feature_segment_id: env_feature_overrides_data.num_segment_overrides += 1 - for identity_override in get_overrides_data_future.result(): - # Only override features that exists in core - if identity_override.feature_state.feature.id in all_overrides_data: - all_overrides_data[ - identity_override.feature_state.feature.id - ].add_identity_override() + for identity_overrides_v2_list in get_overrides_data_future.result(): + + for identity_override in identity_overrides_v2_list.identity_overrides: + # Only override features that exists in core + if identity_override.feature_state.feature.id in all_overrides_data: + all_overrides_data[ + identity_override.feature_state.feature.id + ].add_identity_override() + all_overrides_data[ + identity_override.feature_state.feature.id + ].is_num_identity_overrides_complete = ( + identity_overrides_v2_list.is_num_identity_overrides_complete + ) return all_overrides_data diff --git a/api/features/serializers.py b/api/features/serializers.py index 7c954758f96f..9b6a88693894 100644 --- a/api/features/serializers.py +++ b/api/features/serializers.py @@ -148,6 +148,12 @@ class CreateFeatureSerializer(DeleteBeforeUpdateWritableNestedModelSerializer): "in the environment provided by the `environment` query parameter. " "Note: will return null for Edge enabled projects." ) + is_num_identity_overrides_complete = serializers.SerializerMethodField( + help_text="A boolean that indicates whether there are more" + " identity overrides than are being listed, if `False`. This field is " + "`True` when querying overrides data for a features list page and " + "exact data has been returned." + ) last_modified_in_any_environment = serializers.SerializerMethodField( help_text="Datetime representing the last time that the feature was modified " @@ -181,6 +187,7 @@ class Meta: "environment_feature_state", "num_segment_overrides", "num_identity_overrides", + "is_num_identity_overrides_complete", "is_server_key_only", "last_modified_in_any_environment", "last_modified_in_current_environment", @@ -298,6 +305,14 @@ def get_num_identity_overrides(self, instance) -> typing.Optional[int]: except (KeyError, AttributeError): return None + def get_is_num_identity_overrides_complete(self, instance) -> typing.Optional[int]: + try: + return self.context["overrides_data"][ + instance.id + ].is_num_identity_overrides_complete + except (KeyError, AttributeError): + return None + def get_last_modified_in_any_environment( self, instance: Feature ) -> datetime | None: diff --git a/api/features/views.py b/api/features/views.py index c12084565f38..0b539730d934 100644 --- a/api/features/views.py +++ b/api/features/views.py @@ -165,10 +165,10 @@ def get_queryset(self): if environment_id: page = self.paginate_queryset(queryset) - self.environment = Environment.objects.get(id=environment_id) + self.feature_ids = [feature.id for feature in page] q = Q( - feature_id__in=[feature.id for feature in page], + feature_id__in=self.feature_ids, identity__isnull=True, feature_segment__isnull=True, ) @@ -205,7 +205,6 @@ def perform_destroy(self, instance): def get_serializer_context(self): context = super().get_serializer_context() - feature_states = getattr(self, "_feature_states", {}) project = get_object_or_404(Project.objects.all(), pk=self.kwargs["project_pk"]) context.update( @@ -216,7 +215,9 @@ def get_serializer_context(self): environment = get_object_or_404( Environment, id=self.request.query_params["environment"] ) - context["overrides_data"] = get_overrides_data(environment) + context["overrides_data"] = get_overrides_data( + environment, self.feature_ids + ) return context diff --git a/api/tests/unit/edge_api/identities/test_edge_api_identities_views.py b/api/tests/unit/edge_api/identities/test_edge_api_identities_views.py index 8f0d3c5c1a87..a5ed40442c34 100644 --- a/api/tests/unit/edge_api/identities/test_edge_api_identities_views.py +++ b/api/tests/unit/edge_api/identities/test_edge_api_identities_views.py @@ -159,7 +159,8 @@ def test_get_edge_identity_overrides_for_a_feature( } mock_dynamodb_wrapper.get_identity_overrides_by_environment_id.assert_called_once_with( - environment_id=environment.id, feature_id=feature.id + environment_id=environment.id, + feature_id=feature.id, ) 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 8b5bd7a4d512..2b4c003c248d 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 @@ -59,6 +59,82 @@ 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__set_feature_ids__return_expected( + settings: SettingsWrapper, + environment: Environment, + flagsmith_environments_v2_table: Table, + feature: Feature, +) -> None: + # Given + settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name + wrapper = DynamoEnvironmentV2Wrapper() + + identity_uuid = str(uuid.uuid4()) + identifier = "identity1" + override_document = { + "environment_id": str(environment.id), + "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": {}, + } + + environment_document = map_environment_to_environment_v2_document(environment) + + flagsmith_environments_v2_table.put_item(Item=override_document) + flagsmith_environments_v2_table.put_item(Item=environment_document) + + # When + results = wrapper.get_identity_overrides_by_environment_id( + environment_id=environment.id, + feature_ids=[feature.id], + ) + + # Then + assert len(results) == 1 + assert results[0].items == [override_document] + assert results[0].is_num_identity_overrides_complete is True + + +def test_environment_v2_wrapper__get_identity_overrides_by_environment_id_with_paging__set_feature_ids__return_expected( + settings: SettingsWrapper, + environment: Environment, + flagsmith_environments_v2_table: Table, + feature: Feature, +) -> None: + # Given + settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name + wrapper = DynamoEnvironmentV2Wrapper() + + environment_document = map_environment_to_environment_v2_document(environment) + flagsmith_environments_v2_table.put_item(Item=environment_document) + for i in range(10_000): + identifier = f"identity{i}" + override_document = { + "environment_id": str(environment.id), + "document_key": get_environments_v2_identity_override_document_key( + feature_id=feature.id, identity_uuid=str(uuid.uuid4()) + ), + "environment_api_key": environment.api_key, + "identifier": identifier, + "feature_state": {}, + } + flagsmith_environments_v2_table.put_item(Item=override_document) + + # When + results = wrapper.get_identity_overrides_by_environment_id( + environment_id=environment.id, + feature_ids=[feature.id], + ) + + # Then + assert len(results) == 1 + assert 6000 < len(results[0].items) < 6350 + assert results[0].is_num_identity_overrides_complete is False + + def test_environment_v2_wrapper__get_identity_overrides_by_environment_id__last_evaluated_key__call_expected( flagsmith_environments_v2_table: Table, mocker: MockerFixture, diff --git a/api/tests/unit/features/test_unit_features_features_service.py b/api/tests/unit/features/test_unit_features_features_service.py index f59f94dd8590..55f3b17e850b 100644 --- a/api/tests/unit/features/test_unit_features_features_service.py +++ b/api/tests/unit/features/test_unit_features_features_service.py @@ -61,30 +61,34 @@ def distinct_identity_featurestate( @pytest.mark.parametrize( - "enable_dynamo_db, edge_v2_migration_status, expected_overrides_getter_name, expected_kwargs", + "enable_dynamo_db, edge_v2_migration_status, expected_overrides_getter_name, expected_args, expected_kwargs", [ ( True, EdgeV2MigrationStatus.NOT_STARTED, "get_core_overrides_data", + [], {"skip_identity_overrides": True}, ), ( True, EdgeV2MigrationStatus.IN_PROGRESS, "get_core_overrides_data", + [], {"skip_identity_overrides": True}, ), ( True, EdgeV2MigrationStatus.COMPLETE, "get_edge_overrides_data", + [None], {}, ), ( False, ANY, "get_core_overrides_data", + [], {}, ), ], @@ -95,6 +99,7 @@ def test_feature_get_overrides_data__call_expected( enable_dynamo_db: bool, edge_v2_migration_status: str, expected_overrides_getter_name: str, + expected_args: list[None], expected_kwargs: dict[str, bool], ) -> None: # Given @@ -117,6 +122,7 @@ def test_feature_get_overrides_data__call_expected( # Then mocked_override_getters.pop(expected_overrides_getter_name).assert_called_once_with( environment, + *expected_args, **expected_kwargs, ) [remaining_override_getter_mock] = mocked_override_getters.values() @@ -235,8 +241,13 @@ def test_feature_get_edge_overrides_data( ) edge_identity.save(admin_user) + feature_ids = [ + distinct_identity_featurestate.feature.id, + feature.id, + ] + # When - overrides_data = get_edge_overrides_data(environment) + overrides_data = get_edge_overrides_data(environment, feature_ids) # Then assert overrides_data[feature.id].num_identity_overrides == 1 @@ -284,11 +295,12 @@ def test_get_edge_overrides_data_skips_deleted_features( ) edge_identity.save(admin_user) + feature_ids = [distinct_identity_featurestate.feature.id, feature.id] # Now, delete one of the feature feature.delete() # When - overrides_data = get_edge_overrides_data(environment) + overrides_data = get_edge_overrides_data(environment, feature_ids) # Then - we only have one identity override(for the feature that still exists) assert len(overrides_data) == 1 diff --git a/api/tests/unit/features/test_unit_features_views.py b/api/tests/unit/features/test_unit_features_views.py index 248ecf656d3d..893bf7238e6a 100644 --- a/api/tests/unit/features/test_unit_features_views.py +++ b/api/tests/unit/features/test_unit_features_views.py @@ -34,6 +34,7 @@ from environments.identities.models import Identity from environments.models import Environment, EnvironmentAPIKey from environments.permissions.models import UserEnvironmentPermission +from features.dataclasses import EnvironmentFeatureOverridesData from features.feature_types import MULTIVARIATE from features.models import Feature, FeatureSegment, FeatureState from features.multivariate.models import MultivariateFeatureOption @@ -2062,6 +2063,38 @@ def test_list_features_provides_segment_overrides_for_dynamo_enabled_project( assert response_json["results"][0]["num_identity_overrides"] is None +def test_list_features_calls_get_overrides_data_with_feature_ids( + dynamo_enabled_project: Project, + dynamo_enabled_project_environment_one: Environment, + admin_client_new: APIClient, + mocker: MockerFixture, +) -> None: + # Given + feature = Feature.objects.create( + name="test_feature", project=dynamo_enabled_project + ) + url = "%s?environment=%d" % ( + reverse( + "api-v1:projects:project-features-list", args=[dynamo_enabled_project.id] + ), + dynamo_enabled_project_environment_one.id, + ) + mock_get_overrides_data = mocker.patch("features.views.get_overrides_data") + mock_get_overrides_data.return_value = { + feature.id: EnvironmentFeatureOverridesData() + } + + # When + response = admin_client_new.get(url) + + # Then + assert response.status_code == status.HTTP_200_OK + mock_get_overrides_data.assert_called_once_with( + dynamo_enabled_project_environment_one, + [feature.id], + ) + + def test_create_segment_override_reaching_max_limit( admin_client_new: APIClient, feature: Feature, @@ -2686,6 +2719,12 @@ def test_list_features_n_plus_1_without_rbac( django_assert_num_queries: DjangoAssertNumQueries, environment: Environment, ) -> None: + """ + TODO: When running locally, this test can come up with an extra query. + It should be tested against CI to ensure it passes, but it would + be even better to solve the underlying issue while runnig locally. + See: https://github.com/Flagsmith/flagsmith/issues/4898 + """ _assert_list_feature_n_plus_1( staff_client, project,