feat: Count v2 identity overrides for feature state list view (#3164)

api/edge_api/identities/edge_identity_service.py

@@ -7,9 +7,12 @@
 
 
 def get_edge_identity_overrides(
-    environment_id: int, feature_id: int
+    environment_id: int,
+    feature_id: int | None = None,
 ) -> typing.List[IdentityOverrideV2]:
-    override_items = ddb_environment_v2_wrapper.get_identity_overrides_by_feature_id(
-        environment_id=environment_id, feature_id=feature_id
+    override_items = (
+        ddb_environment_v2_wrapper.get_identity_overrides_by_environment_id(
+            environment_id=environment_id, feature_id=feature_id
+        )
     )
     return [IdentityOverrideV2.parse_obj(item) for item in override_items]

api/environments/dynamodb/dynamodb_wrapper.py

@@ -51,33 +51,41 @@
 class BaseDynamoWrapper:
     table_name: str = None
 
-    def __init__(self):
-        self._table: "Table" | None = None
-        if table_name := self.get_table_name():
-            self._table = boto3.resource(
-                "dynamodb", config=Config(tcp_keepalive=True)
-            ).Table(table_name)
+    def __init__(self) -> None:
+        self._table: typing.Optional["Table"] = None
 
     @property
-    def is_enabled(self) -> bool:
-        return self._table is not None
+    def table(self) -> typing.Optional["Table"]:
+        if not self._table:
+            self._table = self.get_table()
+        return self._table
 
-    def get_table_name(self):
+    def get_table_name(self) -> str:
         return self.table_name
 
+    def get_table(self) -> typing.Optional["Table"]:
+        if table_name := self.get_table_name():
+            return boto3.resource("dynamodb", config=Config(tcp_keepalive=True)).Table(
+                table_name
+            )
+
+    @property
+    def is_enabled(self) -> bool:
+        return self.table is not None
+
 
 class DynamoIdentityWrapper(BaseDynamoWrapper):
     def get_table_name(self) -> str | None:
         return settings.IDENTITIES_TABLE_NAME_DYNAMO
 
     def query_items(self, *args, **kwargs) -> "QueryOutputTableTypeDef":
-        return self._table.query(*args, **kwargs)
+        return self.table.query(*args, **kwargs)
 
     def put_item(self, identity_dict: dict):
-        self._table.put_item(Item=identity_dict)
+        self.table.put_item(Item=identity_dict)
 
     def write_identities(self, identities: Iterable["Identity"]):
-        with self._table.batch_writer() as batch:
+        with self.table.batch_writer() as batch:
             for identity in identities:
                 identity_document = map_identity_to_identity_document(identity)
                 # Since sort keys can not be greater than 1024
@@ -90,10 +98,10 @@ def write_identities(self, identities: Iterable["Identity"]):
                 batch.put_item(Item=identity_document)
 
     def get_item(self, composite_key: str) -> typing.Optional[dict]:
-        return self._table.get_item(Key={"composite_key": composite_key}).get("Item")
+        return self.table.get_item(Key={"composite_key": composite_key}).get("Item")
 
     def delete_item(self, composite_key: str):
-        self._table.delete_item(Key={"composite_key": composite_key})
+        self.table.delete_item(Key={"composite_key": composite_key})
 
     def get_item_from_uuid(self, uuid: str) -> dict:
         filter_expression = Key("identity_uuid").eq(uuid)
@@ -200,15 +208,15 @@ class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper):
     table_name = settings.ENVIRONMENTS_TABLE_NAME_DYNAMO
 
     def write_environments(self, environments: Iterable["Environment"]):
-        with self._table.batch_writer() as writer:
+        with self.table.batch_writer() as writer:
             for environment in environments:
                 writer.put_item(
                     Item=map_environment_to_environment_document(environment),
                 )
 
     def get_item(self, api_key: str) -> dict:
         try:
-            return self._table.get_item(Key={"api_key": api_key})["Item"]
+            return self.table.get_item(Key={"api_key": api_key})["Item"]
         except KeyError as e:
             raise ObjectDoesNotExist() from e
 
@@ -217,13 +225,13 @@ class DynamoEnvironmentV2Wrapper(BaseDynamoEnvironmentWrapper):
     def get_table_name(self) -> str | None:
         return settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO
 
-    def get_identity_overrides_by_feature_id(
+    def get_identity_overrides_by_environment_id(
         self,
         environment_id: int,
-        feature_id: int,
+        feature_id: int | None = None,
     ) -> typing.List[dict[str, Any]]:
         try:
-            response = self._table.query(
+            response = self.table.query(
                 KeyConditionExpression=Key(ENVIRONMENTS_V2_PARTITION_KEY).eq(
                     str(environment_id),
                 )
@@ -246,7 +254,7 @@ def update_identity_overrides(
             changeset.to_delete,
             chunk_size=DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT,
         ):
-            with self._table.batch_writer() as writer:
+            with self.table.batch_writer() as writer:
                 for identity_override_to_delete in to_delete:
                     writer.delete_item(
                         Key={
@@ -262,7 +270,7 @@ def update_identity_overrides(
                     )
 
     def write_environments(self, environments: Iterable["Environment"]) -> None:
-        with self._table.batch_writer() as writer:
+        with self.table.batch_writer() as writer:
             for environment in environments:
                 writer.put_item(
                     Item=map_environment_to_environment_v2_document(environment),
@@ -276,7 +284,7 @@ def write_api_key(self, api_key: "EnvironmentAPIKey"):
         self.write_api_keys([api_key])
 
     def write_api_keys(self, api_keys: Iterable["EnvironmentAPIKey"]):
-        with self._table.batch_writer() as writer:
+        with self.table.batch_writer() as writer:
             for api_key in api_keys:
                 writer.put_item(
                     Item=map_environment_api_key_to_environment_api_key_document(

api/environments/dynamodb/utils.py

@@ -1,23 +1,13 @@
-from multimethod import overload
-
-# TODO This might require type: ignores in the future, but it's just so nice!
-
-
-@overload
-def get_environments_v2_identity_override_document_key() -> str:
-    return "identity_override:"
-
-
-@overload
-def get_environments_v2_identity_override_document_key(  # noqa: F811
-    feature_id: int,
-) -> str:
-    return f"identity_override:{feature_id}:"
-
-
-@overload
-def get_environments_v2_identity_override_document_key(  # noqa: F811
-    feature_id: int,
-    identity_uuid: str,
+def get_environments_v2_identity_override_document_key(
+    feature_id: int | None = None,
+    identity_uuid: str | None = None,
 ) -> str:
+    if feature_id is None:
+        if identity_uuid:
+            raise ValueError(
+                "Cannot generate identity override document key without feature_id"
+            )
+        return "identity_override:"
+    if identity_uuid is None:
+        return f"identity_override:{feature_id}:"
     return f"identity_override:{feature_id}:{identity_uuid}"

api/features/features_service.py

@@ -1,33 +1,107 @@
 import typing
+from concurrent.futures import ThreadPoolExecutor
 
+from edge_api.identities.edge_identity_service import (
+    get_edge_identity_overrides,
+)
 from features.dataclasses import EnvironmentFeatureOverridesData
 from features.versioning.versioning_service import get_environment_flags_list
+from projects.models import IdentityOverridesV2MigrationStatus
 
 if typing.TYPE_CHECKING:
     from environments.models import Environment
 
 
+OverridesData = dict[int, EnvironmentFeatureOverridesData]
+
+
 def get_overrides_data(
     environment: "Environment",
-) -> typing.Dict[int, EnvironmentFeatureOverridesData]:
+) -> OverridesData:
+    """
+    Get correct overrides counts for a given environment.
+
+    :param project: project to get overrides data for
+    :return: overrides data getter
+    """
+    project = environment.project
+    match project.enable_dynamo_db, project.identity_overrides_v2_migration_status:
+        case True, IdentityOverridesV2MigrationStatus.COMPLETE:
+            # If v2 migration is complete, count segment overrides from Core
+            # and identity overrides from DynamoDB.
+            return get_edge_overrides_data(environment)
+        case True, _:
+            # If v2 migration is in progress or not started, we want to count Core overrides,
+            # but only the segment ones, as the identity ones in DynamoDB are uncountable for v1.
+            return get_core_overrides_data(
+                environment,
+                skip_identity_overrides=True,
+            )
+        case _, _:
+            # For projects still fully on Core, count all overrides from Core.
+            return get_core_overrides_data(environment)
+
+
+def get_core_overrides_data(
+    environment: "Environment",
+    *,
+    skip_identity_overrides: bool = False,
+) -> OverridesData:
     """
     Get the number of identity / segment overrides in a given environment for each feature in the
     project.
 
     :param environment: the environment to get the overrides data for
-    :return: dictionary of {feature_id: EnvironmentFeatureOverridesData}
+    :return OverridesData: dictionary of {feature_id: EnvironmentFeatureOverridesData}
     """
     environment_feature_states_list = get_environment_flags_list(environment)
-    all_overrides_data = {}
+    all_overrides_data: OverridesData = {}
 
     for feature_state in environment_feature_states_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
+        elif skip_identity_overrides:
+            continue
         elif feature_state.identity_id:
             env_feature_overrides_data.add_identity_override()
-        all_overrides_data[feature_state.feature_id] = env_feature_overrides_data
+
+    return all_overrides_data
+
+
+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.
+    Retrieve identity override data from DynamoDB.
+
+    :param environment: the environment to get the overrides data for
+    :return OverridesData: dictionary of {feature_id: EnvironmentFeatureOverridesData}
+    """
+    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,
+            environment_id=environment.id,
+        )
+    all_overrides_data: OverridesData = {}
+
+    for feature_state in get_environment_flags_list_future.result():
+        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_override in get_overrides_data_future.result():
+        all_overrides_data[
+            identity_override.feature_state.feature.id
+        ].add_identity_override()
 
     return all_overrides_data

api/features/views.py

@@ -152,13 +152,12 @@ def perform_destroy(self, instance):
 
     def get_serializer_context(self):
         context = super().get_serializer_context()
-        if self.kwargs.get("project_pk"):
-            context.update(
-                project=get_object_or_404(
-                    Project.objects.all(), pk=self.kwargs["project_pk"]
-                ),
-                user=self.request.user,
-            )
+        context.update(
+            project=get_object_or_404(
+                Project.objects.all(), pk=self.kwargs["project_pk"]
+            ),
+            user=self.request.user,
+        )
         if self.action == "list" and "environment" in self.request.query_params:
             environment = get_object_or_404(
                 Environment, id=self.request.query_params["environment"]