feat(api): Evaluate identity flags through flag-engine

api/edge_api/identities/models.py

@@ -6,6 +6,7 @@
 from django.conf import settings
 from django.db.models import Prefetch, Q
 from django.utils import timezone
+from flag_engine.engine import get_evaluation_result
 
 from api_keys.user import APIKeyUser
 from edge_api.identities.tasks import (
@@ -17,6 +18,8 @@
 from edge_api.identities.utils import generate_change_dict
 from environments.dynamodb import DynamoIdentityWrapper
 from environments.models import Environment
+from evaluation.mappers import map_feature_state_to_feature_context
+from evaluation.types import EvaluationContext
 from features.models import FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
 from features.versioning.versioning_service import get_environment_flags_dict
@@ -103,32 +106,59 @@ def get_all_feature_states(
             | Q(feature_segment__isnull=True)
         )
 
-        feature_states: dict[str, FeatureState | FeatureStateModel] = (
-            get_environment_flags_dict(  # type: ignore[assignment]
-                environment=django_environment,
-                additional_filters=additional_filters,
-                additional_select_related_args=[
-                    "feature",
-                    "feature_segment",
-                    "feature_segment__segment",
-                    "feature_state_value",
-                ],
-                additional_prefetch_related_args=[
-                    Prefetch(
-                        "multivariate_feature_state_values",
-                        queryset=MultivariateFeatureStateValue.objects.select_related(
-                            "multivariate_feature_option"
-                        ),
-                    )
-                ],
-                # since we only want to retrieve the highest priority feature state,
-                # we key off the feature name instead of the default
-                # (feature_id, segment_id, identity_id). This will give us only e.g.
-                # the highest priority matching segment override for a given feature.
-                key_function=lambda fs: fs.feature.name,  # type: ignore[arg-type,return-value]
-            )
+        django_feature_states: dict[str, FeatureState] = get_environment_flags_dict(  # type: ignore[assignment]
+            environment=django_environment,
+            additional_filters=additional_filters,
+            additional_select_related_args=[
+                "feature",
+                "feature_segment",
+                "feature_segment__segment",
+                "feature_state_value",
+            ],
+            additional_prefetch_related_args=[
+                Prefetch(
+                    "multivariate_feature_state_values",
+                    queryset=MultivariateFeatureStateValue.objects.select_related(
+                        "multivariate_feature_option"
+                    ),
+                )
+            ],
+            # since we only want to retrieve the highest priority feature state,
+            # we key off the feature name instead of the default
+            # (feature_id, segment_id, identity_id). This will give us only e.g.
+            # the highest priority matching segment override for a given feature.
+            key_function=lambda fs: fs.feature.name,  # type: ignore[arg-type,return-value]
         )
 
+        # The winning row per feature is settled above, but the value an
+        # identity is served for a multivariate feature is the engine's to
+        # decide. Ask it, and carry the answer on each row.
+        context: EvaluationContext = {
+            "environment": {
+                "key": django_environment.api_key,
+                "name": django_environment.name or "",
+            },
+            "identity": {
+                "identifier": self.identifier,
+                "key": self.get_hash_key(
+                    django_environment.use_identity_composite_key_for_hashing
+                ),
+            },
+            "features": {
+                feature_name: map_feature_state_to_feature_context(
+                    feature_state,
+                    mv_fs_values=feature_state.multivariate_feature_state_values.all(),
+                )
+                for feature_name, feature_state in django_feature_states.items()
+            },
+        }
+        for flag in get_evaluation_result(context)["flags"].values():
+            django_feature_states[flag["name"]].flag_result = flag
+
+        feature_states: dict[str, FeatureState | FeatureStateModel] = {
+            **django_feature_states
+        }
+
         # Since the identity overrides are the highest priority, we can now iterate
         # over the dictionary and replace any feature states with those that have
         # an identity override, stored against the identity in dynamo.

api/environments/identities/models.py

@@ -1,18 +1,15 @@
 from itertools import chain
 
 from django.db import models
-from django.db.models import Prefetch, Q
 from flag_engine.engine import get_evaluation_result
 
 from environments.identities.managers import IdentityManager
 from environments.identities.traits.models import Trait
 from environments.models import Environment
 from environments.sdk.types import SDKTraitData
+from evaluation.mappers import map_environment_to_evaluation_context
 from features.models import FeatureState
-from features.multivariate.models import MultivariateFeatureStateValue
-from features.versioning.versioning_service import get_environment_flags_list
 from segments.models import Segment
-from util.mappers.engine import map_environment_to_evaluation_context
 
 
 class Identity(models.Model):
@@ -52,81 +49,6 @@ def get_hash_key(self, use_identity_composite_key_for_hashing: bool = False) ->
             else str(self.id)
         )
 
-    def get_all_feature_states(
-        self,
-        feature_name: str | None = None,
-        traits: list[Trait] | None = None,
-        additional_filters: Q | None = None,
-    ) -> list[FeatureState]:
-        """
-        Get all feature states for an identity. This method returns a single flag for
-        each feature in the identity's environment's project. The flag returned is the
-        correct flag based on the priorities as follows (highest -> lowest):
-
-            1. Identity - flag override for this specific identity
-            2. Segment - flag overridden for a segment this identity belongs to
-            3. Environment - default value for the environment
-
-        :return: (list) flags for an identity with the correct values based on
-            identity / segment priorities
-        """
-        segments = self.get_segments(traits=traits, overrides_only=True)
-
-        # define sub queries
-        belongs_to_environment_query = Q(environment=self.environment)
-        if self.id:
-            overridden_for_identity_query = Q(identity=self)
-        else:
-            # skip identity overrides for transient identities
-            overridden_for_identity_query = Q()
-
-        overridden_for_segment_query = Q(
-            feature_segment__segment__in=segments,
-            feature_segment__environment=self.environment,
-        )
-        environment_default_query = Q(identity=None, feature_segment=None)
-
-        # define the full query
-        full_query = belongs_to_environment_query & (
-            overridden_for_identity_query
-            | overridden_for_segment_query
-            | environment_default_query
-        )
-
-        if additional_filters:
-            full_query &= additional_filters
-
-        all_flags = get_environment_flags_list(
-            environment=self.environment,
-            feature_name=feature_name,
-            additional_filters=full_query,
-            additional_prefetch_related_args=[
-                Prefetch(
-                    "multivariate_feature_state_values",
-                    queryset=MultivariateFeatureStateValue.objects.select_related(
-                        "multivariate_feature_option"
-                    ),
-                )
-            ],
-        )
-
-        # iterate over all the flags and build a dictionary keyed on feature with the highest priority flag
-        # for the given identity as the value.
-        identity_flags = {}
-        for flag in all_flags:
-            if flag.feature_id not in identity_flags:
-                identity_flags[flag.feature_id] = flag
-            else:
-                current_flag = identity_flags[flag.feature_id]
-                if flag > current_flag:
-                    identity_flags[flag.feature_id] = flag
-
-        if self.environment.get_hide_disabled_flags() is True:
-            # filter out any flags that are disabled
-            return [value for value in identity_flags.values() if value.enabled]
-
-        return list(identity_flags.values())
-
     def get_overridden_feature_states(self) -> dict[int, FeatureState]:
         """
         Get all overridden feature states for an identity.
@@ -166,8 +88,11 @@ def get_segments(
         )
         result = get_evaluation_result(context)
         return [
-            segments_by_pk[segment_result["metadata"]["pk"]]
+            segments_by_pk[pk]
             for segment_result in result["segments"]
+            # Synthetic identity-override segments carry no pk, and are not
+            # segments as far as any caller is concerned.
+            if (pk := segment_result["metadata"].get("pk")) is not None
         ]
 
     def get_all_user_traits(self):  # type: ignore[no-untyped-def]

api/environments/identities/serializers.py

@@ -96,20 +96,19 @@ class IdentityAllFeatureStatesSerializer(serializers.Serializer):  # type: ignor
     def get_feature_state_value(
         self, instance: typing.Union[FeatureState, FeatureStateModel]
     ) -> typing.Union[str, int, bool]:
-        identity = self.context["identity"]
-        environment_api_key = self.context["environment_api_key"]
+        if isinstance(instance, FeatureState):
+            return instance.evaluated_value  # type: ignore[no-any-return]
 
-        environment = Environment.get_from_cache(environment_api_key)
+        # An edge identity's own overrides are stored in DynamoDB rather than
+        # the ORM, and are still resolved outside the engine.
+        environment = Environment.get_from_cache(self.context["environment_api_key"])
         assert environment
-        hash_key = identity.get_hash_key(
-            environment.use_identity_composite_key_for_hashing
+        return instance.get_value(  # type: ignore[no-any-return]
+            self.context["identity"].get_hash_key(
+                environment.use_identity_composite_key_for_hashing
+            )
         )
 
-        if isinstance(instance, FeatureState):
-            return instance.get_feature_state_value_by_hash_key(hash_key)  # type: ignore[no-any-return]
-
-        return instance.get_value(hash_key)  # type: ignore[no-any-return]
-
     def get_overridden_by(self, instance) -> typing.Optional[str]:  # type: ignore[no-untyped-def]
         if getattr(instance, "feature_segment_id", None) is not None:
             return "SEGMENT"

api/environments/identities/views.py

@@ -34,6 +34,7 @@
     IdentifyWithTraitsSerializer,
     IdentitySerializerWithTraitsAndSegments,
 )
+from evaluation.services import get_identity_feature_states
 from features.serializers import SDKIdentityFeatureStateSerializer
 from integrations.integration import identify_integrations
 from util.views import SDKAPIView
@@ -134,7 +135,7 @@ def get(self, request, identifier, *args, **kwargs):  # type: ignore[no-untyped-
             "IdentityTraitFlagsSegments", ("flags", "traits", "segments")
         )
         identity_flags_traits_segments = IdentityFlagsWithTraitsAndSegments(
-            flags=identity.get_all_feature_states(),
+            flags=get_identity_feature_states(identity),
             traits=traits_data,
             segments=identity.get_segments(),
         )
@@ -289,7 +290,8 @@ def _get_single_feature_state_response(
     ) -> Response:
         context = self.get_serializer_context()  # type: ignore[no-untyped-call]
 
-        for feature_state in identity.get_all_feature_states(
+        for feature_state in get_identity_feature_states(
+            identity,
             additional_filters=self._get_additional_filters(),
         ):
             if feature_state.feature.name == feature_name:
@@ -317,7 +319,8 @@ def _get_all_feature_states_for_user_response(  # type: ignore[no-untyped-def]
         :param identity: Identity model to return feature states for
         :return: Response containing lists of both serialized flags and traits
         """
-        all_feature_states = identity.get_all_feature_states(
+        all_feature_states = get_identity_feature_states(
+            identity,
             additional_filters=self._get_additional_filters(),
         )
         serializer_class = self.get_serializer_class()

api/environments/sdk/serializers.py

@@ -17,6 +17,7 @@
     get_transient_identity_and_traits,
 )
 from environments.sdk.types import SDKTraitData
+from evaluation.services import get_identity_feature_states
 from features.serializers import (
     FeatureStateSerializerFull,
     SDKIdentityFeatureStateSerializer,
@@ -177,7 +178,8 @@ def save(self, **kwargs):  # type: ignore[no-untyped-def]
                 sdk_trait_data=sdk_trait_data,
             )
 
-        all_feature_states = identity.get_all_feature_states(
+        all_feature_states = get_identity_feature_states(
+            identity,
             traits=traits,
             additional_filters=self.context.get("feature_states_additional_filters"),
         )