fix: Speed up identity overrides (#4840)

Co-authored-by: Matthew Elwell <[email protected]>

api/edge_api/identities/edge_identity_service.py

@@ -1,18 +1,20 @@
-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()
 
 
 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

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

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,

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):
         """

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,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
@@ -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

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:

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