fix: Speed up identity overrides

api/edge_api/identities/edge_identity_service.py

@@ -9,10 +9,13 @@
 def get_edge_identity_overrides(
     environment_id: int,
     feature_id: int | None = None,
+    feature_ids: None | list[int] = None,
 ) -> typing.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,
+            feature_ids=feature_ids,
         )
     )
     return [

api/environments/dynamodb/types.py

@@ -88,6 +88,7 @@ class IdentityOverrideV2(BaseModel):
     identifier: str
     identity_uuid: str
     feature_state: FeatureStateModel
+    more_identity_overrides: bool = False
 
 
 @dataclass

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -1,4 +1,5 @@
 import typing
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Iterable
 
 from boto3.dynamodb.conditions import Key
@@ -66,23 +67,68 @@ def get_identity_overrides_by_environment_id(
         self,
         environment_id: int,
         feature_id: int | None = None,
+        feature_ids: None | list[int] = None,
     ) -> typing.List[dict[str, Any]]:
         try:
-            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(
+            if feature_ids is None:
+                return list(
+                    self.query_get_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 = []
+                for future in futures:
+                    result = future.result()
+                    for item in result:
+                        results.append(item)
+
+                return results
+
         except KeyError as e:
             raise ObjectDoesNotExist() from e
 
+    def get_identity_overrides_page(
+        self, environment_id: int, feature_id: int
+    ) -> list[dict[str, Any]]:
+        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")
+        for item in query_response["Items"]:
+            item["more_identity_overrides"] = last_evaluated_key is not None
+        return query_response["Items"]
+
+    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,

api/features/dataclasses.py

@@ -14,6 +14,7 @@ class EnvironmentFeatureOverridesData:
 
     num_segment_overrides: int = 0
     num_identity_overrides: typing.Optional[int] = None
+    more_identity_overrides: bool = False
 
     def add_identity_override(self):
         """

api/features/features_service.py

@@ -16,20 +16,21 @@
 
 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
 
     if project.enable_dynamo_db:
         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
@@ -89,6 +90,7 @@ def get_edge_overrides_data(
         get_overrides_data_future = executor.submit(
             get_edge_identity_overrides,
             environment_id=environment.id,
+            feature_ids=feature_ids,
         )
     all_overrides_data: OverridesData = {}
 
@@ -104,5 +106,8 @@ def get_edge_overrides_data(
             all_overrides_data[
                 identity_override.feature_state.feature.id
             ].add_identity_override()
+            all_overrides_data[
+                identity_override.feature_state.feature.id
+            ].more_identity_overrides = identity_override.more_identity_overrides
 
     return all_overrides_data

api/features/serializers.py

@@ -148,6 +148,11 @@ class CreateFeatureSerializer(DeleteBeforeUpdateWritableNestedModelSerializer):
         "in the environment provided by the `environment` query parameter. "
         "Note: will return null for Edge enabled projects."
     )
+    more_identity_overrides = serializers.SerializerMethodField(
+        help_text="A boolean that indicates whether there are more"
+        " identity overrides than are being listed in order to "
+        "improve performance of the features list endpoint."
+    )
 
     last_modified_in_any_environment = serializers.SerializerMethodField(
         help_text="Datetime representing the last time that the feature was modified "
@@ -181,6 +186,7 @@ class Meta:
             "environment_feature_state",
             "num_segment_overrides",
             "num_identity_overrides",
+            "more_identity_overrides",
             "is_server_key_only",
             "last_modified_in_any_environment",
             "last_modified_in_current_environment",
@@ -298,6 +304,12 @@ def get_num_identity_overrides(self, instance) -> typing.Optional[int]:
         except (KeyError, AttributeError):
             return None
 
+    def get_more_identity_overrides(self, instance) -> typing.Optional[int]:
+        try:
+            return self.context["overrides_data"][instance.id].more_identity_overrides
+        except (KeyError, AttributeError):
+            return None
+
     def get_last_modified_in_any_environment(
         self, instance: Feature
     ) -> datetime | None:

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
 

api/tests/unit/edge_api/identities/test_edge_api_identities_views.py

@@ -159,7 +159,9 @@ 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,
+        feature_ids=None,
     )
 
 

api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py

@@ -59,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__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
+    override_document["more_identity_overrides"] = False
+    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,

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()

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,10 @@ def test_list_features_n_plus_1_without_rbac(
     django_assert_num_queries: DjangoAssertNumQueries,
     environment: Environment,
 ) -> None:
+    """
+    NOTE: When running locally, this test can come up with an extra query.
+          It should be tested against CI to ensure it passes.
+    """
     _assert_list_feature_n_plus_1(
         staff_client,
         project,