fix: Speed up identity overrides

api/edge_api/identities/edge_identity_service.py

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

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -1,4 +1,5 @@
 import typing
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Iterable
 
 from boto3.dynamodb.conditions import Key
@@ -66,23 +67,70 @@ def get_identity_overrides_by_environment_id(
         self,
         environment_id: int,
         feature_id: int | None = None,
+        feature_ids: None | list[int] = None,
+        limit_feature_identities_to_one_page: bool = False,
     ) -> typing.List[dict[str, Any]]:
         try:
-            return list(
-                self.query_get_all_items(
-                    KeyConditionExpression=Key(ENVIRONMENTS_V2_PARTITION_KEY).eq(
-                        str(environment_id),
+            if not limit_feature_identities_to_one_page:
+                if feature_ids is not None:
+                    raise NotImplementedError(
+                        "Multiple feature ids is currently not supported "
+                        "when not limiting to a single page of features"
                     )
-                    & 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,
+                            ),
+                        )
                     )
                 )
-            )
+            else:
+                if feature_ids is None and feature_id:
+                    feature_ids = [feature_id]
+                assert feature_ids is not None
+
+                futures = []
+                with ThreadPoolExecutor() as executor:
+                    for feature_id in feature_ids:
+                        futures.append(
+                            executor.submit(
+                                self.get_page_of_feature_identities,
+                                environment_id,
+                                feature_id,
+                            )
+                        )
+
+                results = []
+                for future in as_completed(futures):
+                    result = future.result()
+                    for item in result:
+                        results.append(item)
+
+                return results
+
         except KeyError as e:
             raise ObjectDoesNotExist() from e
 
+    def get_page_of_feature_identities(
+        self, environment_id: int, feature_id: int
+    ) -> list[dict[str, Any]]:
+        query_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 query_response["Items"]
+
     def update_identity_overrides(
         self,
         changeset: IdentityOverridesV2Changeset,

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,8 @@ def get_edge_overrides_data(
         get_overrides_data_future = executor.submit(
             get_edge_identity_overrides,
             environment_id=environment.id,
+            feature_ids=feature_ids,
+            limit_feature_identities_to_one_page=True,
         )
     all_overrides_data: OverridesData = {}
 

api/features/views.py

@@ -165,7 +165,6 @@ def get_queryset(self):
 
         if environment_id:
             page = self.paginate_queryset(queryset)
-
             self.environment = Environment.objects.get(id=environment_id)
             q = Q(
                 feature_id__in=[feature.id for feature in page],
@@ -205,7 +204,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 +214,10 @@ def get_serializer_context(self):
             environment = get_object_or_404(
                 Environment, id=self.request.query_params["environment"]
             )
-            context["overrides_data"] = get_overrides_data(environment)
+            queryset = self.get_queryset()
+            page = self.paginate_queryset(queryset)
+            feature_ids = [feature.id for feature in page]
+            context["overrides_data"] = get_overrides_data(environment, feature_ids)
 
         return context
 

api/tests/unit/features/test_unit_features_views.py

@@ -33,6 +33,7 @@
     VIEW_ENVIRONMENT,
 )
 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,