fix: Rely on Flagsmith Engine for segment evaluation, avoid N+1 queri…

…es (#3038)

api/conftest.py

@@ -3,6 +3,7 @@
 import pytest
 from django.contrib.contenttypes.models import ContentType
 from django.core.cache import cache
+from flag_engine.segments.constants import EQUAL
 from rest_framework.authtoken.models import Token
 from rest_framework.test import APIClient
 
@@ -46,7 +47,7 @@
 )
 from projects.permissions import VIEW_PROJECT
 from projects.tags.models import Tag
-from segments.models import EQUAL, Condition, Segment, SegmentRule
+from segments.models import Condition, Segment, SegmentRule
 from task_processor.task_run_method import TaskRunMethod
 from users.models import FFAdminUser, UserPermissionGroup
 

api/environments/dynamodb/tests/test_dynamodb_identity_wrapper.py

@@ -5,12 +5,13 @@
 from core.constants import INTEGER
 from django.core.exceptions import ObjectDoesNotExist
 from flag_engine.identities.builders import build_identity_model
+from flag_engine.segments.constants import IN
 from rest_framework.exceptions import NotFound
 
 from environments.dynamodb import DynamoIdentityWrapper
 from environments.identities.models import Identity
 from environments.identities.traits.models import Trait
-from segments.models import IN, Condition, Segment, SegmentRule
+from segments.models import Condition, Segment, SegmentRule
 from util.mappers import (
     map_environment_to_environment_document,
     map_identity_to_identity_document,

api/environments/identities/managers.py

@@ -1,6 +1,32 @@
+from typing import TYPE_CHECKING, Iterable
+
 from django.db.models import Manager
 
+if TYPE_CHECKING:
+    from environments.identities.models import Identity
+    from environments.models import Environment
+    from integrations.integration import IntegrationConfig
+
 
-class IdentityManager(Manager):
+class IdentityManager(Manager["Identity"]):
     def get_by_natural_key(self, identifier, environment_api_key):
         return self.get(identifier=identifier, environment__api_key=environment_api_key)
+
+    def get_or_create_for_sdk(
+        self,
+        identifier: str,
+        environment: "Environment",
+        integrations: Iterable["IntegrationConfig"],
+    ) -> tuple["Identity", bool]:
+        return (
+            self.select_related(
+                "environment",
+                "environment__project",
+                *[
+                    f"environment__{integration['relation_name']}"
+                    for integration in integrations
+                ],
+            )
+            .prefetch_related("identity_traits")
+            .get_or_create(identifier=identifier, environment=environment)
+        )

api/environments/identities/models.py

@@ -3,13 +3,19 @@
 from django.db import models
 from django.db.models import Prefetch, Q
 from django.utils import timezone
+from flag_engine.segments.evaluator import evaluate_identity_in_segment
 
 from environments.identities.managers import IdentityManager
 from environments.identities.traits.models import Trait
 from environments.models import Environment
 from features.models import FeatureState
 from features.multivariate.models import MultivariateFeatureStateValue
 from segments.models import Segment
+from util.mappers.engine import (
+    map_identity_to_engine,
+    map_segment_to_engine,
+    map_traits_to_engine,
+)
 
 
 class Identity(models.Model):
@@ -153,8 +159,21 @@ def get_segments(
         else:
             all_segments = self.environment.project.get_segments_from_cache()
 
+        engine_identity = map_identity_to_engine(
+            self,
+            with_overrides=False,
+            with_traits=False,
+        )
+        engine_traits = map_traits_to_engine(traits)
+
         for segment in all_segments:
-            if segment.does_identity_match(self, traits=traits):
+            engine_segment = map_segment_to_engine(segment)
+
+            if evaluate_identity_in_segment(
+                identity=engine_identity,
+                segment=engine_segment,
+                override_traits=engine_traits,
+            ):
                 matching_segments.append(segment)
 
         return matching_segments

api/environments/identities/tests/test_models.py

@@ -1,6 +1,13 @@
 import pytest
 from core.constants import FLOAT
 from django.utils import timezone
+from flag_engine.segments.constants import (
+    EQUAL,
+    GREATER_THAN,
+    GREATER_THAN_INCLUSIVE,
+    LESS_THAN_INCLUSIVE,
+    NOT_EQUAL,
+)
 from rest_framework.test import APITestCase
 
 from environments.identities.models import Identity
@@ -15,16 +22,7 @@
 from features.value_types import BOOLEAN, INTEGER, STRING
 from organisations.models import Organisation
 from projects.models import Project
-from segments.models import (
-    EQUAL,
-    GREATER_THAN,
-    GREATER_THAN_INCLUSIVE,
-    LESS_THAN_INCLUSIVE,
-    NOT_EQUAL,
-    Condition,
-    Segment,
-    SegmentRule,
-)
+from segments.models import Condition, Segment, SegmentRule
 
 from .helpers import (
     create_trait_for_identity,

api/environments/identities/tests/test_views.py

@@ -4,10 +4,11 @@
 from unittest.case import TestCase
 
 import pytest
-from core.constants import FLAGSMITH_UPDATED_AT_HEADER
+from core.constants import FLAGSMITH_UPDATED_AT_HEADER, STRING
 from django.test import override_settings
 from django.urls import reverse
 from django.utils import timezone
+from flag_engine.segments.constants import PERCENTAGE_SPLIT
 from rest_framework import status
 from rest_framework.test import APIClient, APITestCase
 
@@ -21,7 +22,6 @@
 from integrations.amplitude.models import AmplitudeConfiguration
 from organisations.models import Organisation, OrganisationRole
 from projects.models import Project
-from segments import models
 from segments.models import Condition, Segment, SegmentRule
 from util.tests import Helper
 
@@ -371,7 +371,7 @@ def test_identities_endpoint_returns_traits(self, mock_amplitude_wrapper):
         trait = Trait.objects.create(
             identity=self.identity,
             trait_key="trait_key",
-            value_type="STRING",
+            value_type=STRING,
             string_value="trait_value",
         )
 
@@ -423,7 +423,7 @@ def test_identities_endpoint_returns_value_for_segment_if_identity_in_segment(
         Trait.objects.create(
             identity=self.identity,
             trait_key=trait_key,
-            value_type="STRING",
+            value_type=STRING,
             string_value=trait_value,
         )
         segment = Segment.objects.create(name="Test Segment", project=self.project)
@@ -477,7 +477,7 @@ def test_identities_endpoint_returns_value_for_segment_if_identity_in_segment_an
         Trait.objects.create(
             identity=self.identity,
             trait_key=trait_key,
-            value_type="STRING",
+            value_type=STRING,
             string_value=trait_value,
         )
         segment = Segment.objects.create(name="Test Segment", project=self.project)
@@ -529,7 +529,7 @@ def test_identities_endpoint_returns_value_for_segment_if_rule_type_percentage_s
             [segment.id, self.identity.id]
         )
         Condition.objects.create(
-            operator=models.PERCENTAGE_SPLIT,
+            operator=PERCENTAGE_SPLIT,
             value=(identity_percentage_value + (1 - identity_percentage_value) / 2)
             * 100.0,
             rule=segment_rule,
@@ -576,7 +576,7 @@ def test_identities_endpoint_returns_default_value_if_rule_type_percentage_split
             [segment.id, self.identity.id]
         )
         Condition.objects.create(
-            operator=models.PERCENTAGE_SPLIT,
+            operator=PERCENTAGE_SPLIT,
             value=identity_percentage_value / 2,
             rule=segment_rule,
         )
@@ -629,13 +629,13 @@ def test_post_identify_deletes_a_trait_if_trait_value_is_none(self):
         trait_1 = Trait.objects.create(
             identity=self.identity,
             trait_key="trait_key_1",
-            value_type="STRING",
+            value_type=STRING,
             string_value="trait_value",
         )
         trait_2 = Trait.objects.create(
             identity=self.identity,
             trait_key="trait_key_2",
-            value_type="STRING",
+            value_type=STRING,
             string_value="trait_value",
         )
 

api/environments/identities/views.py

@@ -112,11 +112,11 @@ class SDKIdentitiesDeprecated(SDKAPIView):
     def get(self, request, identifier, *args, **kwargs):
         # if we have identifier fetch, or create if does not exist
         if identifier:
-            identity, _ = Identity.objects.get_or_create(
+            identity, _ = Identity.objects.get_or_create_for_sdk(
                 identifier=identifier,
                 environment=request.environment,
+                integrations=IDENTITY_INTEGRATIONS,
             )
-
         else:
             return Response(
                 {"detail": "Missing identifier"}, status=status.HTTP_400_BAD_REQUEST
@@ -172,17 +172,10 @@ def get(self, request):
                 {"detail": "Missing identifier"}
             )  # TODO: add 400 status - will this break the clients?
 
-        identity, _ = (
-            Identity.objects.select_related(
-                "environment",
-                "environment__project",
-                *[
-                    f"environment__{integration['relation_name']}"
-                    for integration in IDENTITY_INTEGRATIONS
-                ],
-            )
-            .prefetch_related("identity_traits")
-            .get_or_create(identifier=identifier, environment=request.environment)
+        identity, _ = Identity.objects.get_or_create_for_sdk(
+            identifier=identifier,
+            environment=request.environment,
+            integrations=IDENTITY_INTEGRATIONS,
         )
         self.identity = identity
 

api/features/feature_segments/tests/test_models.py

@@ -1,14 +1,15 @@
 import pytest
 from core.constants import STRING
 from django.test import TestCase
+from flag_engine.segments.constants import EQUAL
 
 from environments.identities.models import Identity
 from environments.identities.traits.models import Trait
 from environments.models import Environment
 from features.models import Feature, FeatureSegment
 from organisations.models import Organisation
 from projects.models import Project
-from segments.models import EQUAL, Condition, Segment, SegmentRule
+from segments.models import Condition, Segment, SegmentRule
 
 
 @pytest.mark.django_db

api/integrations/integration.py

@@ -1,11 +1,20 @@
+from typing import Type, TypedDict
+
 from integrations.amplitude.amplitude import AmplitudeWrapper
+from integrations.common.wrapper import AbstractBaseIdentityIntegrationWrapper
 from integrations.heap.heap import HeapWrapper
 from integrations.mixpanel.mixpanel import MixpanelWrapper
 from integrations.rudderstack.rudderstack import RudderstackWrapper
 from integrations.segment.segment import SegmentWrapper
 from integrations.webhook.webhook import WebhookWrapper
 
-IDENTITY_INTEGRATIONS = [
+
+class IntegrationConfig(TypedDict):
+    relation_name: str
+    wrapper: Type[AbstractBaseIdentityIntegrationWrapper]
+
+
+IDENTITY_INTEGRATIONS: list[IntegrationConfig] = [
     {"relation_name": "amplitude_config", "wrapper": AmplitudeWrapper},
     {"relation_name": "segment_config", "wrapper": SegmentWrapper},
     {"relation_name": "heap_config", "wrapper": HeapWrapper},

api/integrations/webhook/serializers.py

@@ -1,13 +1,15 @@
 import typing
 
 from django.db.models import Q
+from flag_engine.segments.evaluator import evaluate_identity_in_segment
 from rest_framework import serializers
 
 from features.serializers import FeatureStateSerializerFull
 from integrations.common.serializers import (
     BaseEnvironmentIntegrationModelSerializer,
 )
 from segments.models import Segment
+from util.mappers.engine import map_identity_to_engine, map_segment_to_engine
 
 from .models import WebhookConfiguration
 
@@ -25,8 +27,16 @@ class Meta:
         model = Segment
         fields = ("id", "name", "member")
 
-    def get_member(self, obj):
-        return obj.does_identity_match(identity=self.context.get("identity"))
+    def get_member(self, obj: Segment) -> bool:
+        engine_identity = map_identity_to_engine(
+            self.context.get("identity"),
+            with_overrides=False,
+        )
+        engine_segment = map_segment_to_engine(obj)
+        return evaluate_identity_in_segment(
+            identity=engine_identity,
+            segment=engine_segment,
+        )
 
 
 class IntegrationFeatureStateSerializer(FeatureStateSerializerFull):