fixes

api/environments/identities/managers.py

@@ -1,30 +1,30 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterable
 
 from django.db.models import Manager
 
-from integrations.integration import IDENTITY_INTEGRATIONS
-
 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 IDENTITY_INTEGRATIONS
+                    for integration in integrations
                 ],
             )
             .prefetch_related("identity_traits")

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):
@@ -152,8 +158,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/views.py

@@ -31,7 +31,10 @@
     IdentitySerializerWithTraitsAndSegments,
 )
 from features.serializers import SDKFeatureStateSerializer
-from integrations.integration import identify_integrations
+from integrations.integration import (
+    IDENTITY_INTEGRATIONS,
+    identify_integrations,
+)
 from util.views import SDKAPIView
 
 
@@ -112,6 +115,7 @@ def get(self, request, identifier, *args, **kwargs):
             identity, _ = Identity.objects.get_or_create_for_sdk(
                 identifier=identifier,
                 environment=request.environment,
+                integrations=IDENTITY_INTEGRATIONS,
             )
         else:
             return Response(
@@ -171,6 +175,7 @@ def get(self, request):
         identity, _ = Identity.objects.get_or_create_for_sdk(
             identifier=identifier,
             environment=request.environment,
+            integrations=IDENTITY_INTEGRATIONS,
         )
         self.identity = identity
 

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):

api/segments/models.py

@@ -10,22 +10,11 @@
 from django.core.exceptions import ValidationError
 from django.db import models
 from flag_engine.segments import constants
-from flag_engine.segments.evaluator import evaluate_identity_in_segment
 
 from audit.constants import SEGMENT_CREATED_MESSAGE, SEGMENT_UPDATED_MESSAGE
 from audit.related_object_type import RelatedObjectType
 from features.models import Feature
 from projects.models import Project
-from util.mappers.engine import (
-    map_identity_to_engine,
-    map_segment_to_engine,
-    map_traits_to_trait_models,
-)
-
-if typing.TYPE_CHECKING:
-    from environments.identities.models import Identity
-    from environments.identities.traits.models import Trait
-
 
 logger = logging.getLogger(__name__)
 
@@ -83,23 +72,6 @@ def id_exists_in_rules_data(rules_data: typing.List[dict]) -> bool:
 
         return False
 
-    def does_identity_match(
-        self, identity: "Identity", traits: typing.List["Trait"] = None
-    ) -> bool:
-        segment_model = map_segment_to_engine(self)
-        identity_model = map_identity_to_engine(
-            identity,
-            with_overrides=False,
-            with_traits=not traits,
-        )
-        trait_models = map_traits_to_trait_models(traits) if traits else None
-
-        return evaluate_identity_in_segment(
-            identity=identity_model,
-            segment=segment_model,
-            override_traits=trait_models,
-        )
-
     def get_create_log_message(self, history_instance) -> typing.Optional[str]:
         return SEGMENT_CREATED_MESSAGE % self.name
 

api/util/mappers/engine.py

@@ -45,10 +45,12 @@
     "map_feature_to_engine",
     "map_identity_to_engine",
     "map_mv_option_to_engine",
+    "map_segment_to_engine",
+    "map_traits_to_engine",
 )
 
 
-def map_traits_to_trait_models(traits: Iterable["Trait"]) -> list[TraitModel]:
+def map_traits_to_engine(traits: Iterable["Trait"]) -> list[TraitModel]:
     return [
         TraitModel(trait_key=trait.trait_key, trait_value=trait.trait_value)
         for trait in traits
@@ -386,7 +388,7 @@ def map_identity_to_engine(
         )
         for feature_state in identity_feature_states
     ]
-    identity_trait_models = map_traits_to_trait_models(identity_traits)
+    identity_trait_models = map_traits_to_engine(identity_traits)
 
     return IdentityModel(
         # Attributes: