fix: metrics endpoint identity overrides (#6159)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Zaimwa9

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -199,34 +199,3 @@ def get_segment_ids(
             return [segment.id for segment in segments]
 
         return []
-
-    def get_identity_overrides_count(self, environment_api_key: str) -> int:
-        return sum(
-            len({f["feature_id"] for f in identity["identity_features"]})
-            for identity in self.iter_all_items_paginated(
-                environment_api_key=environment_api_key,
-                overrides_only=True,
-            )
-        )
-
-    def get_identity_override_feature_counts(
-        self, environment_api_key: str
-    ) -> dict[int, int]:
-        feature_to_identity_count: dict[int, int] = {}
-
-        for identity in self.iter_all_items_paginated(
-            environment_api_key=environment_api_key,
-            overrides_only=True,
-        ):
-            unique_feature_ids: set[int] = set()
-
-            for feature_override in identity.get("identity_features", []):
-                feature_id = feature_override.get("feature", {}).get("id", 0)
-                unique_feature_ids.add(feature_id)
-
-            for feature_id in unique_feature_ids:
-                feature_to_identity_count[feature_id] = (
-                    feature_to_identity_count.get(feature_id, 0) + 1
-                )
-
-        return feature_to_identity_count

api/environments/models.py

@@ -479,11 +479,12 @@ def get_scheduled_metrics_queryset(self) -> QuerySet[FeatureState]:
                 & Q(change_sets__live_from__gt=timezone.now())
             )
         )
-
-        change_requests: QuerySet["ChangeRequest"] = qs.filter(
-            id__in=qs.filter(scheduled_q).values_list("id", flat=True)
+        ids = qs.filter(scheduled_q).values_list("id", flat=True)
+        unique_ids = list(set(ids))
+        change_requests_qs: QuerySet["ChangeRequest"] = ChangeRequest.objects.filter(
+            id__in=unique_ids
         )
-        return change_requests
+        return change_requests_qs
 
     @staticmethod
     def is_bad_key(environment_key: str) -> bool:

api/metrics/metrics_service.py

@@ -5,7 +5,7 @@
 from django.db import models
 from django.db.models import Q
 
-from edge_api.identities.models import EdgeIdentity
+from edge_api.identities import edge_identity_service
 from environments.models import Environment
 from features.models import Feature
 from metrics.constants import DEFAULT_METRIC_DEFINITIONS, WORKFLOW_METRIC_DEFINITIONS
@@ -83,26 +83,26 @@ def _get_identity_metrics(self) -> dict[EnvMetricsName, Callable[[], int]]:
         }
 
     def _get_active_identity_edge_overrides_count(self) -> int:
-        override_feature_id_counts = (
-            EdgeIdentity.dynamo_wrapper.get_identity_override_feature_counts(
-                self.environment.api_key
-            )
-        )
-
-        valid_feature_ids = set(
+        environment_feature_ids = set(
             Feature.objects.filter(
                 project=self.environment.project,
                 is_archived=False,
                 deleted_at__isnull=True,
             ).values_list("id", flat=True)
         )
 
-        return sum(
-            count
-            for feature_id, count in override_feature_id_counts.items()
-            if feature_id in valid_feature_ids
+        all_overrides = edge_identity_service.get_edge_identity_overrides(
+            environment_id=self.environment.id,
+            feature_id=None,
         )
 
+        count = 0
+        for override in all_overrides:
+            if override.feature_state.feature.id in environment_feature_ids:
+                count += 1
+
+        return count
+
     def _get_feature_metrics(
         self,
     ) -> dict[EnvMetricsName, Callable[[], int]]:

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

@@ -601,74 +601,3 @@ def test_delete_all_identities__deletes_all_identities_documents_from_dynamodb(
     # Then
     assert flagsmith_identities_table.scan()["Count"] == 1
     assert flagsmith_identities_table.scan()["Items"][0] == identity_three
-
-
-@pytest.mark.parametrize(
-    "identity_features, expected_counts",
-    [
-        (
-            [[1, 2, 3], [99], []],
-            {1: 1, 2: 1, 3: 1, 99: 1},
-        ),
-        (
-            [[1, 2, 3], [99], [1, 2, 3, 99]],
-            {1: 2, 2: 2, 3: 2, 99: 2},
-        ),
-        (
-            [[], [], []],
-            {},
-        ),
-        (
-            [[], [1, 2, 3], []],
-            {1: 1, 2: 1, 3: 1},
-        ),
-        (
-            [[4], [4, 25, 18, 19, 85, 100], [4]],
-            {4: 3, 25: 1, 18: 1, 19: 1, 85: 1, 100: 1},
-        ),
-    ],
-)
-def test_get_identity_override_feature_counts_dynamo_returns_correct_total(
-    flagsmith_identities_table: Table,
-    dynamodb_identity_wrapper: DynamoIdentityWrapper,
-    identity_features: list[list[int]],
-    expected_counts: dict[int, int],
-) -> None:
-    environment_api_key = "env_test"
-
-    identity_one = {
-        "composite_key": f"{environment_api_key}_identity1",
-        "environment_api_key": environment_api_key,
-        "identifier": "user1",
-        "identity_features": [
-            {"feature": {"id": feature_id}} for feature_id in identity_features[0]
-        ],
-    }
-
-    identity_two = {
-        "composite_key": f"{environment_api_key}_identity2",
-        "environment_api_key": environment_api_key,
-        "identifier": "user2",
-        "identity_features": [
-            {"feature": {"id": feature_id}} for feature_id in identity_features[1]
-        ],
-    }
-
-    identity_three = {
-        "composite_key": f"{environment_api_key}_identity3",
-        "environment_api_key": environment_api_key,
-        "identifier": "user3",
-        "identity_features": [
-            {"feature": {"id": feature_id}} for feature_id in identity_features[2]
-        ],
-    }
-
-    flagsmith_identities_table.put_item(Item=identity_one)
-    flagsmith_identities_table.put_item(Item=identity_two)
-    flagsmith_identities_table.put_item(Item=identity_three)
-
-    result = dynamodb_identity_wrapper.get_identity_override_feature_counts(
-        environment_api_key
-    )
-
-    assert result == expected_counts

api/tests/unit/metrics/test_unit_metrics_service.py

@@ -72,6 +72,7 @@ def test_environment_metrics_service_builds_expected_metrics(
 
 
 @pytest.mark.parametrize("uses_dynamo, expected_value", [(True, 99), (False, 1)])
+@pytest.mark.django_db
 def test_dynamo_identity_metric_used(
     monkeypatch: pytest.MonkeyPatch,
     environment: Environment,
@@ -83,8 +84,7 @@ def test_dynamo_identity_metric_used(
     monkeypatch.setattr(
         environment, "get_segment_metrics_queryset", lambda: MagicMock(count=lambda: 1)
     )
-    Feature.objects.create(
-        id=10,
+    feature = Feature.objects.create(
         project=environment.project,
         name="feature-10",
         is_archived=False,
@@ -97,14 +97,20 @@ def test_dynamo_identity_metric_used(
         lambda: MagicMock(count=identity_count_mock),
     )
 
-    dynamo_mock = MagicMock(return_value={10: 99})
+    mock_override = MagicMock()
+    mock_override.feature_state.feature.id = feature.id
+    mock_overrides = [mock_override] * 99
+
+    dynamo_mock = MagicMock(return_value=mock_overrides)
     monkeypatch.setattr(
-        "edge_api.identities.models.EdgeIdentity.dynamo_wrapper.get_identity_override_feature_counts",
+        "edge_api.identities.edge_identity_service.get_edge_identity_overrides",
         dynamo_mock,
     )
+
     # When
     metrics_service = EnvironmentMetricsService(environment)
     metrics = metrics_service.get_metrics_payload()
+
     # Then
     assert metrics_service.uses_dynamo == uses_dynamo
     assert any(
@@ -114,7 +120,10 @@ def test_dynamo_identity_metric_used(
     )
 
     if uses_dynamo:
-        dynamo_mock.assert_called_once()
+        dynamo_mock.assert_called_once_with(
+            environment_id=environment.id,
+            feature_id=None,
+        )
         identity_count_mock.assert_not_called()
     else:
         identity_count_mock.assert_called_once()