fix: N dynamodb queries in feature list endpoint (#6758)

Author: matthewelwell

api/edge_api/identities/edge_identity_service.py

@@ -3,7 +3,6 @@
 from edge_api.identities.models import EdgeIdentity
 from environments.dynamodb import DynamoEnvironmentV2Wrapper
 from environments.dynamodb.types import (
-    IdentityOverridesV2List,
     IdentityOverrideV2,
 )
 
@@ -22,40 +21,24 @@ def get_edge_identity_overrides(
     )
     return [
         IdentityOverrideV2.model_validate(
-            {**item, "environment_id": str(item["environment_id"])}  # type: ignore[dict-item,index]
+            {**item, "environment_id": str(item["environment_id"])}
         )
         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 = (
+def get_edge_identity_override_keys(environment_id: int) -> list[str]:
+    """
+    Get all the identity overrides for an environment, returning only the document key
+    for optimised performance when the key is all that is needed.
+    """
+    override_items = (
         ddb_environment_v2_wrapper.get_identity_overrides_by_environment_id(
             environment_id=environment_id,
-            feature_ids=feature_ids,
+            projection_expression_attributes=["document_key"],
         )
     )
-
-    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  # type: ignore[union-attr]
-        ]
-        complete = identity_overrides_query_response.is_num_identity_overrides_complete  # type: ignore[union-attr]
-        results.append(
-            IdentityOverridesV2List(
-                identity_overrides=identity_overrides,
-                is_num_identity_overrides_complete=complete,
-            )
-        )
-
-    return results
+    return [item["document_key"] for item in override_items]
 
 
 def get_overridden_feature_ids_for_edge_identity(identity_uuid: str) -> set[int]:

api/environments/dynamodb/types.py

@@ -97,12 +97,6 @@ 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/utils.py

@@ -11,3 +11,7 @@ def get_environments_v2_identity_override_document_key(
     if identity_uuid is None:
         return f"identity_override:{feature_id}:"
     return f"identity_override:{feature_id}:{identity_uuid}"
+
+
+def get_feature_id_from_identity_override_document_key(document_key: str) -> int:
+    return int(document_key.split(":")[1])

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -1,6 +1,4 @@
 import typing
-from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
 from typing import Any, Iterable
 
 from boto3.dynamodb.conditions import Key
@@ -31,12 +29,6 @@
     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])
@@ -74,52 +66,25 @@ def get_identity_overrides_by_environment_id(
         self,
         environment_id: int,
         feature_id: int | None = None,
-        feature_ids: None | list[int] = None,
-    ) -> list[dict[str, Any]] | list[IdentityOverridesQueryResponse]:
-        try:
-            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
+        projection_expression_attributes: list[str] | None = None,
+    ) -> list[dict[str, Any]]:
+        key_condition_expression = self.get_identity_overrides_key_condition_expression(
+            environment_id=environment_id,
+            feature_id=feature_id,
+        )
+        query_kwargs: dict[str, Any] = {
+            "KeyConditionExpression": key_condition_expression,
+        }
+        if projection_expression_attributes:
+            query_kwargs["ProjectionExpression"] = ",".join(
+                projection_expression_attributes
+            )
 
+        try:
+            return list(self.query_iter_all_items(**query_kwargs))
         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(  # type: ignore[union-attr]
-            KeyConditionExpression=self.get_identity_overrides_key_condition_expression(  # type: ignore[arg-type]
-                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,

api/features/dataclasses.py

@@ -14,7 +14,6 @@ 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):  # type: ignore[no-untyped-def]
         """

api/features/features_service.py

@@ -2,7 +2,10 @@
 from concurrent.futures import ThreadPoolExecutor
 
 from edge_api.identities.edge_identity_service import (
-    get_edge_identity_overrides_for_feature_ids,
+    get_edge_identity_override_keys,
+)
+from environments.dynamodb.utils import (
+    get_feature_id_from_identity_override_document_key,
 )
 from features.dataclasses import EnvironmentFeatureOverridesData
 from features.versioning.versioning_service import get_environment_flags_list
@@ -16,12 +19,11 @@
 
 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
+    :param environment: environment to get overrides data for
     :return: overrides data getter dictionary of {feature_id: EnvironmentFeatureOverridesData}
     """
     project = environment.project
@@ -30,7 +32,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, feature_ids)
+            return get_edge_overrides_data(environment)
         # If v2 migration is not started, in progress, or incomplete,
         # only count segment overrides from Core.
         # v1 Edge identity overrides are uncountable.
@@ -71,9 +73,7 @@ def get_core_overrides_data(
     return all_overrides_data
 
 
-def get_edge_overrides_data(
-    environment: "Environment", feature_ids: None | list[int] = None
-) -> OverridesData:
+def get_edge_overrides_data(environment: "Environment") -> OverridesData:
     """
     Get the number of identity / segment overrides in a given environment for each feature in the
     project.
@@ -83,37 +83,32 @@ def get_edge_overrides_data(
     :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,
-        )
+        # Note: We intentionally let the get_environment_flags_list
+        # call happen on the main thread to simplify testing.
+        # The call to dynamo happens on a separate thread, which
+        # still gives us the parallelism we need without needing
+        # an extra thread. This does mean that the order of execution
+        # is important here.
         get_overrides_data_future = executor.submit(
-            get_edge_identity_overrides_for_feature_ids,
+            get_edge_identity_override_keys,
             environment_id=environment.id,
-            feature_ids=feature_ids,
         )
+        flags_list = get_environment_flags_list(environment)
     all_overrides_data: OverridesData = {}
 
-    for feature_state in get_environment_flags_list_future.result():
+    for feature_state in flags_list:
         env_feature_overrides_data = all_overrides_data.setdefault(
             feature_state.feature_id, EnvironmentFeatureOverridesData()
         )
         if feature_state.feature_segment_id:
             env_feature_overrides_data.num_segment_overrides += 1
-    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[  # type: ignore[no-untyped-call]
-                    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
-                )
+    for identity_override_key in get_overrides_data_future.result():
+        feature_id = get_feature_id_from_identity_override_document_key(
+            identity_override_key
+        )
+        # Only override features that exists in core
+        if feature_id in all_overrides_data:
+            all_overrides_data[feature_id].add_identity_override()  # type: ignore[no-untyped-call]
 
     return all_overrides_data

api/features/serializers.py

@@ -180,11 +180,10 @@ 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."
+
+    # This is kept for backwards compatibility, but is always true
+    is_num_identity_overrides_complete = serializers.BooleanField(
+        read_only=True, default=True
     )
 
     last_modified_in_any_environment = serializers.SerializerMethodField(
@@ -345,14 +344,6 @@ def get_num_identity_overrides(self, instance: Feature) -> int | None:
         except (KeyError, AttributeError):
             return None
 
-    def get_is_num_identity_overrides_complete(self, instance: Feature) -> bool | None:
-        try:
-            return self.context["overrides_data"][  # type: ignore[no-any-return]
-                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

@@ -349,9 +349,7 @@ def get_serializer_context(self):  # type: ignore[no-untyped-def]
             environment = get_object_or_404(
                 Environment, id=self.request.query_params["environment"]
             )
-            context["overrides_data"] = get_overrides_data(
-                environment, self.feature_ids
-            )
+            context["overrides_data"] = get_overrides_data(environment)
 
         return context
 

api/tests/conftest.py

@@ -1,4 +1,5 @@
 import typing
+import uuid
 
 import pytest
 from django.conf import settings
@@ -11,11 +12,18 @@
     DynamoEnvironmentWrapper,
     DynamoIdentityWrapper,
 )
+from environments.dynamodb.types import 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 tests.types import MigratorFactory
 from util.mappers import (
     map_environment_to_environment_document,
     map_environment_to_environment_v2_document,
+    map_feature_state_to_engine,
+    map_identity_override_to_identity_override_document,
 )
 
 
@@ -102,6 +110,39 @@ def dynamo_enabled_project_environment_one_document(
     return environment_dict
 
 
+@pytest.fixture()
+def identity_override_document(
+    flagsmith_environments_v2_table: Table,
+    environment: Environment,
+    feature: Feature,
+) -> dict[str, typing.Any]:
+    identity_uuid = str(uuid.uuid4())
+    identifier = "identity-with-dynamo-override"
+
+    identity_override = IdentityOverrideV2(
+        environment_id=str(environment.id),
+        environment_api_key=environment.api_key,
+        document_key=get_environments_v2_identity_override_document_key(
+            feature_id=feature.id, identity_uuid=identity_uuid
+        ),
+        identifier=identifier,
+        identity_uuid=identity_uuid,
+        feature_state=map_feature_state_to_engine(
+            FeatureState(
+                feature=feature,
+                enabled=True,
+                environment=environment,
+            ),
+        ),
+    )
+
+    identity_override_document = map_identity_override_to_identity_override_document(
+        identity_override
+    )
+    flagsmith_environments_v2_table.put_item(Item=identity_override_document)
+    return identity_override_document
+
+
 @pytest.fixture()
 def dynamo_environment_wrapper(
     flagsmith_environment_table: Table,

api/tests/unit/conftest.py

@@ -10,7 +10,7 @@
 from environments.models import Environment
 from features.models import Feature
 from organisations.models import Organisation, OrganisationRole
-from projects.models import Project
+from projects.models import EdgeV2MigrationStatus, Project
 from projects.tags.models import Tag
 from segments.models import Segment
 from users.models import FFAdminUser
@@ -112,6 +112,7 @@ def dynamo_enabled_project(organisation):  # type: ignore[no-untyped-def]
         name="Dynamo enabled project",
         organisation=organisation,
         enable_dynamo_db=True,
+        edge_v2_migration_status=EdgeV2MigrationStatus.COMPLETE,
     )