feat: search identities by dashboard alias (#4569)

api/conftest.py

@@ -12,6 +12,7 @@
 from flag_engine.segments.constants import EQUAL
 from moto import mock_dynamodb
 from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table
+from pytest_django.fixtures import SettingsWrapper
 from pytest_django.plugin import blocking_manager_key
 from pytest_mock import MockerFixture
 from rest_framework.authtoken.models import Token
@@ -977,8 +978,10 @@ def dynamodb(aws_credentials):
 
 
 @pytest.fixture()
-def flagsmith_identities_table(dynamodb: DynamoDBServiceResource) -> Table:
-    return dynamodb.create_table(
+def flagsmith_identities_table(
+    dynamodb: DynamoDBServiceResource, settings: SettingsWrapper
+) -> Table:
+    table = dynamodb.create_table(
         TableName="flagsmith_identities",
         KeySchema=[
             {
@@ -991,6 +994,7 @@ def flagsmith_identities_table(dynamodb: DynamoDBServiceResource) -> Table:
             {"AttributeName": "environment_api_key", "AttributeType": "S"},
             {"AttributeName": "identifier", "AttributeType": "S"},
             {"AttributeName": "identity_uuid", "AttributeType": "S"},
+            {"AttributeName": "dashboard_alias", "AttributeType": "S"},
         ],
         GlobalSecondaryIndexes=[
             {
@@ -1006,9 +1010,24 @@ def flagsmith_identities_table(dynamodb: DynamoDBServiceResource) -> Table:
                 "KeySchema": [{"AttributeName": "identity_uuid", "KeyType": "HASH"}],
                 "Projection": {"ProjectionType": "ALL"},
             },
+            {
+                "IndexName": "environment_api_key-dashboard_alias-index",
+                "KeySchema": [
+                    {"AttributeName": "environment_api_key", "KeyType": "HASH"},
+                    {"AttributeName": "dashboard_alias", "KeyType": "RANGE"},
+                ],
+                "Projection": {
+                    "ProjectionType": "INCLUDE",
+                    "NonKeyAttributes": [
+                        "identifier",
+                    ],
+                },
+            },
         ],
         BillingMode="PAY_PER_REQUEST",
     )
+    settings.IDENTITIES_TABLE_NAME_DYNAMO = table.name
+    return table
 
 
 @pytest.fixture()

api/edge_api/identities/search.py

@@ -0,0 +1,28 @@
+import enum
+from dataclasses import dataclass
+
+IDENTIFIER_ATTRIBUTE = "identifier"
+DASHBOARD_ALIAS_ATTRIBUTE = "dashboard_alias"
+DASHBOARD_ALIAS_SEARCH_PREFIX = f"{DASHBOARD_ALIAS_ATTRIBUTE}:"
+
+
+class EdgeIdentitySearchType(enum.Enum):
+    EQUAL = "EQUAL"
+    BEGINS_WITH = "BEGINS_WITH"
+
+
+@dataclass
+class EdgeIdentitySearchData:
+    search_term: str
+    search_type: EdgeIdentitySearchType
+    search_attribute: str
+
+    @property
+    def dynamo_search_method(self):
+        return (
+            "eq" if self.search_type == EdgeIdentitySearchType.EQUAL else "begins_with"
+        )
+
+    @property
+    def dynamo_index_name(self):
+        return f"environment_api_key-{self.search_attribute}-index"

api/edge_api/identities/serializers.py

@@ -32,18 +32,29 @@
 from webhooks.constants import WEBHOOK_DATETIME_FORMAT
 
 from .models import EdgeIdentity
+from .search import (
+    DASHBOARD_ALIAS_ATTRIBUTE,
+    DASHBOARD_ALIAS_SEARCH_PREFIX,
+    IDENTIFIER_ATTRIBUTE,
+    EdgeIdentitySearchData,
+    EdgeIdentitySearchType,
+)
 from .tasks import call_environment_webhook_for_feature_state_change
 
 
 class EdgeIdentitySerializer(serializers.Serializer):
     identity_uuid = serializers.CharField(read_only=True)
     identifier = serializers.CharField(required=True, max_length=2000)
+    dashboard_alias = serializers.CharField(required=False, max_length=100)
 
-    def save(self, **kwargs):
+    def create(self, *args, **kwargs):
         identifier = self.validated_data.get("identifier")
+        dashboard_alias = self.validated_data.get("dashboard_alias")
         environment_api_key = self.context["view"].kwargs["environment_api_key"]
         self.instance = EngineIdentity(
-            identifier=identifier, environment_api_key=environment_api_key
+            identifier=identifier,
+            environment_api_key=environment_api_key,
+            dashboard_alias=dashboard_alias,
         )
         if EdgeIdentity.dynamo_wrapper.get_item(self.instance.composite_key):
             raise ValidationError(
@@ -55,6 +66,28 @@ def save(self, **kwargs):
         return self.instance
 
 
+class EdgeIdentityUpdateSerializer(EdgeIdentitySerializer):
+    def get_fields(self):
+        fields = super().get_fields()
+        fields["identifier"].read_only = True
+        return fields
+
+    def update(
+        self, instance: dict[str, typing.Any], validated_data: dict[str, typing.Any]
+    ) -> EngineIdentity:
+        engine_identity = EngineIdentity.model_validate(instance)
+
+        engine_identity.dashboard_alias = (
+            self.validated_data.get("dashboard_alias")
+            or engine_identity.dashboard_alias
+        )
+
+        edge_identity = EdgeIdentity(engine_identity)
+        edge_identity.save()
+
+        return engine_identity
+
+
 class EdgeMultivariateFeatureOptionField(serializers.IntegerField):
     def to_internal_value(
         self,
@@ -238,6 +271,36 @@ class GetEdgeIdentityOverridesQuerySerializer(serializers.Serializer):
     feature = serializers.IntegerField(required=False)
 
 
+class EdgeIdentitySearchField(serializers.CharField):
+    def to_internal_value(self, data: str) -> EdgeIdentitySearchData:
+        kwargs = {}
+        search_term = data
+
+        if search_term.startswith(DASHBOARD_ALIAS_SEARCH_PREFIX):
+            kwargs["search_attribute"] = DASHBOARD_ALIAS_ATTRIBUTE
+            search_term = search_term.lstrip(DASHBOARD_ALIAS_SEARCH_PREFIX)
+        else:
+            kwargs["search_attribute"] = IDENTIFIER_ATTRIBUTE
+
+        if search_term.startswith('"') and search_term.endswith('"'):
+            kwargs["search_type"] = EdgeIdentitySearchType.EQUAL
+            search_term = search_term[1:-1]
+        else:
+            kwargs["search_type"] = EdgeIdentitySearchType.BEGINS_WITH
+
+        return EdgeIdentitySearchData(**kwargs, search_term=search_term)
+
+
+class ListEdgeIdentitiesQuerySerializer(serializers.Serializer):
+    page_size = serializers.IntegerField(required=False)
+    q = EdgeIdentitySearchField(
+        required=False,
+        help_text="Search string to look for. Prefix with 'dashboard_alias:' "
+        "to search over aliases instead of identifiers.",
+    )
+    last_evaluated_key = serializers.CharField(required=False, allow_null=True)
+
+
 class GetEdgeIdentityOverridesResultSerializer(serializers.Serializer):
     identifier = serializers.CharField()
     identity_uuid = serializers.CharField()

api/edge_api/identities/views.py

@@ -3,7 +3,6 @@
 import typing
 
 import pydantic
-from boto3.dynamodb.conditions import Key
 from django.shortcuts import get_object_or_404
 from django.utils.decorators import method_decorator
 from drf_yasg.utils import swagger_auto_schema
@@ -22,6 +21,7 @@
     DestroyModelMixin,
     ListModelMixin,
     RetrieveModelMixin,
+    UpdateModelMixin,
 )
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.request import Request
@@ -40,10 +40,12 @@
     EdgeIdentitySerializer,
     EdgeIdentitySourceIdentityRequestSerializer,
     EdgeIdentityTraitsSerializer,
+    EdgeIdentityUpdateSerializer,
     EdgeIdentityWithIdentifierFeatureStateDeleteRequestBody,
     EdgeIdentityWithIdentifierFeatureStateRequestBody,
     GetEdgeIdentityOverridesQuerySerializer,
     GetEdgeIdentityOverridesSerializer,
+    ListEdgeIdentitiesQuerySerializer,
 )
 from environments.identities.serializers import (
     IdentityAllFeatureStatesSerializer,
@@ -66,6 +68,7 @@
     EdgeIdentityWithIdentifierViewPermissions,
     GetEdgeIdentityOverridesPermission,
 )
+from .search import EdgeIdentitySearchData
 
 
 @method_decorator(
@@ -81,14 +84,10 @@ class EdgeIdentityViewSet(
     RetrieveModelMixin,
     DestroyModelMixin,
     ListModelMixin,
+    UpdateModelMixin,
 ):
-    serializer_class = EdgeIdentitySerializer
     pagination_class = EdgeIdentityPagination
     lookup_field = "identity_uuid"
-    dynamo_identifier_search_functions = {
-        "EQUAL": lambda identifier: Key("identifier").eq(identifier),
-        "BEGINS_WITH": lambda identifier: Key("identifier").begins_with(identifier),
-    }
 
     def initial(self, request, *args, **kwargs):
         environment = self.get_environment_from_request()
@@ -97,44 +96,43 @@ def initial(self, request, *args, **kwargs):
 
         super().initial(request, *args, **kwargs)
 
-    def _get_search_function_and_value(
-        self,
-        search_query: str,
-    ) -> typing.Tuple[typing.Callable, str]:
-        if search_query.startswith('"') and search_query.endswith('"'):
-            return self.dynamo_identifier_search_functions[
-                "EQUAL"
-            ], search_query.replace('"', "")
-        return self.dynamo_identifier_search_functions["BEGINS_WITH"], search_query
+    def get_serializer_class(self):
+        if self.action in ("update", "partial_update"):
+            return EdgeIdentityUpdateSerializer
+        return EdgeIdentitySerializer
 
     def get_object(self):
+        # TODO: should this return an EdgeIdentity object instead of a dict?
         return EdgeIdentity.dynamo_wrapper.get_item_from_uuid_or_404(
             self.kwargs["identity_uuid"]
         )
 
     def get_queryset(self):
         page_size = self.pagination_class().get_page_size(self.request)
-        previous_last_evaluated_key = self.request.GET.get("last_evaluated_key")
-        search_query = self.request.query_params.get("q")
+
+        query_serializer = ListEdgeIdentitiesQuerySerializer(
+            data=self.request.query_params
+        )
+        query_serializer.is_valid(raise_exception=True)
+
         start_key = None
-        if previous_last_evaluated_key:
+        if previous_last_evaluated_key := query_serializer.validated_data.get(
+            "last_evaluated_key"
+        ):
             start_key = json.loads(base64.b64decode(previous_last_evaluated_key))
 
-        if not search_query:
+        search_query: typing.Optional[EdgeIdentitySearchData]
+        if not (search_query := query_serializer.validated_data.get("q")):
             return EdgeIdentity.dynamo_wrapper.get_all_items(
                 self.kwargs["environment_api_key"], page_size, start_key
             )
-        search_func, search_identifier = self._get_search_function_and_value(
-            search_query
-        )
-        identity_documents = EdgeIdentity.dynamo_wrapper.search_items_with_identifier(
-            self.kwargs["environment_api_key"],
-            search_identifier,
-            search_func,
-            page_size,
-            start_key,
+
+        return EdgeIdentity.dynamo_wrapper.search_items(
+            environment_api_key=self.kwargs["environment_api_key"],
+            search_data=search_query,
+            limit=page_size,
+            start_key=start_key,
         )
-        return identity_documents
 
     def get_permissions(self):
         return [

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -12,6 +12,7 @@
 from flag_engine.segments.evaluator import get_identity_segments
 from rest_framework.exceptions import NotFound
 
+from edge_api.identities.search import EdgeIdentitySearchData
 from environments.dynamodb.constants import IDENTITIES_PAGINATION_LIMIT
 from environments.dynamodb.wrappers.exceptions import CapacityBudgetExceeded
 from util.mappers import map_identity_to_identity_document
@@ -147,24 +148,29 @@ def iter_all_items_paginated(
             if last_evaluated_key := query_response.get("LastEvaluatedKey"):
                 get_all_items_kwargs["start_key"] = last_evaluated_key
 
-    def search_items_with_identifier(
+    def search_items(
         self,
         environment_api_key: str,
-        identifier: str,
-        search_function: typing.Callable,
+        search_data: EdgeIdentitySearchData,
         limit: int,
         start_key: dict = None,
-    ):
-        filter_expression = Key("environment_api_key").eq(
+    ) -> "QueryOutputTableTypeDef":
+        partition_key_search_expression = Key("environment_api_key").eq(
             environment_api_key
-        ) & search_function(identifier)
+        )
+        sort_key_search_expression = getattr(
+            Key(search_data.search_attribute), search_data.dynamo_search_method
+        )(search_data.search_term)
+
         query_kwargs = {
-            "IndexName": "environment_api_key-identifier-index",
+            "IndexName": search_data.dynamo_index_name,
             "Limit": limit,
-            "KeyConditionExpression": filter_expression,
+            "KeyConditionExpression": partition_key_search_expression
+            & sort_key_search_expression,
         }
         if start_key:
             query_kwargs.update(ExclusiveStartKey=start_key)
+
         return self.query_items(**query_kwargs)
 
     def get_segment_ids(