fix(Admin API): Sort features by overrides

api/edge_api/identities/edge_identity_service.py

@@ -1,3 +1,6 @@
+from django.core.exceptions import ObjectDoesNotExist
+
+from edge_api.identities.models import EdgeIdentity
 from environments.dynamodb import DynamoEnvironmentV2Wrapper
 from environments.dynamodb.types import (
     IdentityOverridesV2List,
@@ -53,3 +56,14 @@ def get_edge_identity_overrides_for_feature_ids(
         )
 
     return results
+
+
+def get_overridden_feature_ids_for_edge_identity(identity_uuid: str) -> set[int]:
+    try:
+        identity_document = EdgeIdentity.dynamo_wrapper.get_item_from_uuid(
+            identity_uuid
+        )
+    except ObjectDoesNotExist:
+        return set()
+    identity = EdgeIdentity.from_identity_document(identity_document)
+    return {fs.feature.id for fs in identity.feature_overrides}

api/features/serializers.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 from typing import Any
+from uuid import UUID
 
 import django.core.exceptions
 from common.features.multivariate.serializers import (
@@ -95,6 +96,11 @@ class FeatureQuerySerializer(serializers.Serializer):  # type: ignore[type-arg]
         required=False,
         help_text="Integer ID of the segment to retrieve segment overrides for.",
     )
+    identity = serializers.CharField(
+        required=False,
+        help_text="ID of the identity to sort features with identity overrides first.",
+    )
+
     is_enabled = serializers.BooleanField(
         allow_null=True,
         required=False,
@@ -116,6 +122,23 @@ class FeatureQuerySerializer(serializers.Serializer):  # type: ignore[type-arg]
         help_text="Comma separated list of group owner ids to filter on",
     )
 
+    @property
+    def project(self) -> Project:
+        if isinstance(project := self.context.get("project"), Project):
+            return project
+        else:  # pragma: no cover
+            raise RuntimeError(f"{type(self)} requires 'project' in context.")
+
+    def validate_identity(self, value: str) -> str:
+        if self.project.enable_dynamo_db:
+            try:
+                UUID(value)
+            except ValueError:
+                raise serializers.ValidationError("Must be a valid UUID.")
+        elif not value.isdigit():
+            raise serializers.ValidationError("Must be a valid integer.")
+        return value
+
     def validate_owners(self, owners: str) -> list[int]:
         try:
             return [int(owner_id.strip()) for owner_id in owners.split(",")]

api/features/views.py

@@ -7,7 +7,17 @@
 from common.projects.permissions import VIEW_PROJECT
 from django.conf import settings
 from django.core.cache import caches
-from django.db.models import Max, Q, QuerySet
+from django.db.models import (
+    BooleanField,
+    Case,
+    Exists,
+    Max,
+    OuterRef,
+    Q,
+    QuerySet,
+    Value,
+    When,
+)
 from django.utils import timezone
 from django.utils.decorators import method_decorator
 from django.views.decorators.cache import cache_page
@@ -35,6 +45,9 @@
 from app_analytics.throttles import InfluxQueryThrottle
 from core.constants import FLAGSMITH_UPDATED_AT_HEADER, SDK_ENVIRONMENT_KEY_HEADER
 from core.request_origin import RequestOrigin
+from edge_api.identities.edge_identity_service import (
+    get_overridden_feature_ids_for_edge_identity,
+)
 from environments.authentication import EnvironmentKeyAuthentication
 from environments.identities.models import Identity
 from environments.identities.serializers import (
@@ -56,7 +69,7 @@
 
 from .constants import INTERSECTION, UNION
 from .features_service import get_overrides_data
-from .models import Feature, FeatureState
+from .models import Feature, FeatureSegment, FeatureState
 from .multivariate.serializers import (
     FeatureMVOptionsValuesResponseSerializer,
 )
@@ -197,15 +210,18 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
             )
         )
 
-        query_serializer = FeatureQuerySerializer(data=self.request.query_params)
+        query_serializer = FeatureQuerySerializer(
+            data=self.request.query_params,
+            context={"project": project},
+        )
         query_serializer.is_valid(raise_exception=True)
         query_data = query_serializer.validated_data
 
         queryset = annotate_feature_queryset_with_code_references_summary(
             queryset, project.id
         )
 
-        queryset = self._filter_queryset(queryset)
+        queryset = self._filter_queryset(queryset, query_serializer)
 
         if environment_id := query_data.get("environment"):
             queryset = queryset.annotate(
@@ -224,7 +240,42 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
             "-" if query_data["sort_direction"] == "DESC" else "",
             query_data["sort_field"],
         )
-        queryset = queryset.order_by(sort)
+        override_ordering: list[str] = []
+        if environment_id and (segment_id := query_data.get("segment")):
+            queryset = queryset.annotate(
+                has_segment_override=Exists(
+                    FeatureSegment.objects.filter(
+                        feature=OuterRef("pk"),
+                        segment_id=segment_id,
+                        environment_id=environment_id,
+                    )
+                ),
+            )
+            override_ordering.append("-has_segment_override")
+        if identity := query_data.get("identity"):
+            if project.enable_dynamo_db:
+                # Bounded by Project.max_features_allowed
+                override_feature_ids = get_overridden_feature_ids_for_edge_identity(
+                    identity
+                )
+                queryset = queryset.annotate(
+                    has_identity_override=Case(
+                        When(pk__in=override_feature_ids, then=Value(True)),
+                        default=Value(False),
+                        output_field=BooleanField(),
+                    ),
+                )
+            else:
+                queryset = queryset.annotate(
+                    has_identity_override=Exists(
+                        FeatureState.objects.filter(
+                            feature=OuterRef("pk"),
+                            identity_id=identity,
+                        )
+                    ),
+                )
+            override_ordering.append("-has_identity_override")
+        queryset = queryset.order_by(*override_ordering, sort)
 
         if environment_id:
             page = self.paginate_queryset(queryset)
@@ -509,9 +560,11 @@ def filter_owners_and_group_owners(
 
         return queryset.filter(owners_q | group_owners_q)
 
-    def _filter_queryset(self, queryset: QuerySet[Feature]) -> QuerySet[Feature]:
-        query_serializer = FeatureQuerySerializer(data=self.request.query_params)
-        query_serializer.is_valid(raise_exception=True)
+    def _filter_queryset(
+        self,
+        queryset: QuerySet[Feature],
+        query_serializer: FeatureQuerySerializer,
+    ) -> QuerySet[Feature]:
         query_data = query_serializer.validated_data
 
         if query_data.get("search"):

api/tests/unit/edge_api/identities/test_edge_identity_service.py

@@ -0,0 +1,116 @@
+from unittest.mock import MagicMock
+
+import pytest
+from django.core.exceptions import ObjectDoesNotExist
+from pytest_mock import MockerFixture
+
+from edge_api.identities.edge_identity_service import (
+    get_overridden_feature_ids_for_edge_identity,
+)
+
+
+@pytest.fixture()
+def _mock_ddb_identity_wrapper(mocker: MockerFixture) -> MagicMock:
+    return mocker.patch(
+        "edge_api.identities.models.EdgeIdentity.dynamo_wrapper",
+    )
+
+
+def _make_identity_document(
+    *,
+    environment_api_key: str,
+    identity_features: list[dict],  # type: ignore[type-arg]
+) -> dict:  # type: ignore[type-arg]
+    return {
+        "composite_key": f"{environment_api_key}_test_user",
+        "identity_traits": [],
+        "identity_features": identity_features,
+        "identifier": "test_user",
+        "created_date": "2021-09-21T10:12:42.230257+00:00",
+        "environment_api_key": environment_api_key,
+        "identity_uuid": "59efa2a7-6a45-46d6-b953-a7073a90eacf",
+        "django_id": None,
+    }
+
+
+def test_get_overridden_feature_ids_for_edge_identity__identity_with_overrides__returns_feature_ids(
+    _mock_ddb_identity_wrapper: MagicMock,
+) -> None:
+    # Given
+    feature_a_id = 101
+    feature_b_id = 202
+    identity_document = _make_identity_document(
+        environment_api_key="test_key",
+        identity_features=[
+            {
+                "feature": {
+                    "id": feature_a_id,
+                    "name": "feature_a",
+                    "type": "STANDARD",
+                },
+                "enabled": True,
+                "featurestate_uuid": "a7495917-ee57-41d1-a64e-e0697dbc57fb",
+                "feature_state_value": None,
+                "feature_segment": None,
+                "multivariate_feature_state_values": [],
+            },
+            {
+                "feature": {
+                    "id": feature_b_id,
+                    "name": "feature_b",
+                    "type": "STANDARD",
+                },
+                "enabled": False,
+                "featurestate_uuid": "b8506028-ff68-42e2-b75f-f1708ecd68fc",
+                "feature_state_value": None,
+                "feature_segment": None,
+                "multivariate_feature_state_values": [],
+            },
+        ],
+    )
+    _mock_ddb_identity_wrapper.get_item_from_uuid.return_value = identity_document
+
+    # When
+    result = get_overridden_feature_ids_for_edge_identity(
+        "59efa2a7-6a45-46d6-b953-a7073a90eacf"
+    )
+
+    # Then
+    assert result == {feature_a_id, feature_b_id}
+    _mock_ddb_identity_wrapper.get_item_from_uuid.assert_called_once_with(
+        "59efa2a7-6a45-46d6-b953-a7073a90eacf"
+    )
+
+
+def test_get_overridden_feature_ids_for_edge_identity__identity_without_overrides__returns_empty_set(
+    _mock_ddb_identity_wrapper: MagicMock,
+) -> None:
+    # Given
+    identity_document = _make_identity_document(
+        environment_api_key="test_key",
+        identity_features=[],
+    )
+    _mock_ddb_identity_wrapper.get_item_from_uuid.return_value = identity_document
+
+    # When
+    result = get_overridden_feature_ids_for_edge_identity(
+        "59efa2a7-6a45-46d6-b953-a7073a90eacf"
+    )
+
+    # Then
+    assert result == set()
+
+
+def test_get_overridden_feature_ids_for_edge_identity__nonexistent_identity__returns_empty_set(
+    _mock_ddb_identity_wrapper: MagicMock,
+) -> None:
+    # Given
+    _mock_ddb_identity_wrapper.get_item_from_uuid.side_effect = ObjectDoesNotExist()
+
+    # When
+    result = get_overridden_feature_ids_for_edge_identity(
+        "59efa2a7-6a45-46d6-b953-a7073a90eacf"
+    )
+
+    # Then
+    assert result == set()