feat: Improve the filter builder

- Use INNER JOIN for identity overrides
- Don't fetch integrations for the SDK endpoint

Author: khvn26

api/environments/constants.py

@@ -0,0 +1,8 @@
+IDENTITY_INTEGRATIONS_RELATION_NAMES = [
+    "amplitude_config",
+    "heap_config",
+    "mixpanel_config",
+    "rudderstack_config",
+    "segment_config",
+    "webhook_config",
+]

api/environments/identities/managers.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Iterable
 
-from django.db.models import Count, Manager, QuerySet
+from django.db.models import Manager, Prefetch, QuerySet
 
 if TYPE_CHECKING:
     from environments.identities.models import Identity
@@ -31,10 +31,10 @@ def get_or_create_for_sdk(
             .get_or_create(identifier=identifier, environment=environment)
         )
 
-    def with_overrides(self) -> QuerySet["Identity"]:
+    def only_overrides(self) -> QuerySet["Identity"]:
         """
         Only include identities with overridden features.
         """
-        return self.annotate(overridden_count=Count("identity_features")).filter(
-            overridden_count__gt=0
-        )
+        return self.prefetch_related(
+            Prefetch("identity_features"),
+        ).filter(identity_features__isnull=False)

api/environments/managers.py

@@ -6,19 +6,19 @@
 
 
 class EnvironmentManager(SoftDeleteManager):
-    def filter_for_document_builder(self, *args, **kwargs):
+    def filter_for_document_builder(
+        self,
+        *args,
+        extra_select_related: list[str] | None = None,
+        extra_prefetch_related: list[Prefetch | str] | None = None,
+        **kwargs,
+    ):
         return (
             super()
             .select_related(
                 "project",
                 "project__organisation",
-                "amplitude_config",
-                "dynatrace_config",
-                "heap_config",
-                "mixpanel_config",
-                "rudderstack_config",
-                "segment_config",
-                "webhook_config",
+                *extra_select_related or (),
             )
             .prefetch_related(
                 Prefetch(
@@ -55,6 +55,7 @@ def filter_for_document_builder(self, *args, **kwargs):
                         "multivariate_feature_option"
                     ),
                 ),
+                *extra_prefetch_related or (),
             )
             .filter(*args, **kwargs)
         )

api/environments/models.py

@@ -1,17 +1,15 @@
-# -*- coding: utf-8 -*-
-from __future__ import unicode_literals
-
 import logging
 import typing
 from copy import deepcopy
 
 from core.models import abstract_base_auditable_model_factory
 from core.request_origin import RequestOrigin
+from django.apps import apps
 from django.conf import settings
 from django.contrib.contenttypes.fields import GenericRelation
 from django.core.cache import caches
 from django.db import models
-from django.db.models import Q
+from django.db.models import Prefetch, Q
 from django.utils import timezone
 from django.utils.translation import ugettext_lazy as _
 from django_lifecycle import (
@@ -36,6 +34,7 @@
     generate_client_api_key,
     generate_server_api_key,
 )
+from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES
 from environments.dynamodb import (
     DynamoEnvironmentAPIKeyWrapper,
     DynamoEnvironmentV2Wrapper,
@@ -213,11 +212,7 @@ def get_from_cache(cls, api_key):
                 select_related_args = (
                     "project",
                     "project__organisation",
-                    "mixpanel_config",
-                    "segment_config",
-                    "amplitude_config",
-                    "heap_config",
-                    "dynatrace_config",
+                    *IDENTITY_INTEGRATIONS_RELATION_NAMES,
                 )
                 base_qs = cls.objects.select_related(*select_related_args).defer(
                     "description"
@@ -243,7 +238,10 @@ def write_environments_to_dynamodb(
             Q(id=environment_id) if environment_id else Q(project_id=project_id)
         )
         environments = list(
-            cls.objects.filter_for_document_builder(environments_filter)
+            cls.objects.filter_for_document_builder(
+                environments_filter,
+                extra_select_related=IDENTITY_INTEGRATIONS_RELATION_NAMES,
+            )
         )
         if not environments:
             return
@@ -369,8 +367,21 @@ def _get_environment_document_from_db(
         cls,
         api_key: str,
     ) -> dict[str, typing.Any]:
-        environment = cls.objects.filter_for_document_builder(api_key=api_key).get()
-        return map_environment_to_sdk_document(environment)
+        Identity = apps.get_model("identities", "Identity")
+        environment = cls.objects.filter_for_document_builder(
+            api_key=api_key,
+            extra_prefetch_related=[
+                Prefetch(
+                    "identities",
+                    queryset=Identity.objects.only_overrides(),
+                    to_attr="identities_with_overrides",
+                )
+            ],
+        ).get()
+        return map_environment_to_sdk_document(
+            environment,
+            identities_with_overrides=environment.identities_with_overrides,
+        )
 
     def _get_environment(self):
         return self

api/integrations/integration.py

@@ -1,5 +1,6 @@
 from typing import Type, TypedDict
 
+from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES
 from integrations.amplitude.amplitude import AmplitudeWrapper
 from integrations.common.wrapper import AbstractBaseIdentityIntegrationWrapper
 from integrations.heap.heap import HeapWrapper
@@ -16,13 +17,24 @@ class IntegrationConfig(TypedDict):
 
 IDENTITY_INTEGRATIONS: list[IntegrationConfig] = [
     {"relation_name": "amplitude_config", "wrapper": AmplitudeWrapper},
-    {"relation_name": "segment_config", "wrapper": SegmentWrapper},
     {"relation_name": "heap_config", "wrapper": HeapWrapper},
     {"relation_name": "mixpanel_config", "wrapper": MixpanelWrapper},
-    {"relation_name": "webhook_config", "wrapper": WebhookWrapper},
     {"relation_name": "rudderstack_config", "wrapper": RudderstackWrapper},
+    {"relation_name": "segment_config", "wrapper": SegmentWrapper},
+    {"relation_name": "webhook_config", "wrapper": WebhookWrapper},
 ]
 
+assert set(IDENTITY_INTEGRATIONS_RELATION_NAMES) == (
+    _configured_integrations := {
+        integration_config["relation_name"]
+        for integration_config in IDENTITY_INTEGRATIONS
+    }
+), (
+    "Check that `environments.constants.IDENTITY_INTEGRATIONS_RELATION_NAMES` and "
+    "`integration.integration.IDENTITY_INTEGRATIONS` contain the same values. \n"
+    f"Unconfigured integrations: {set(IDENTITY_INTEGRATIONS_RELATION_NAMES) - _configured_integrations}"
+)
+
 
 def identify_integrations(identity, all_feature_states, trait_models=None):
     for integration in IDENTITY_INTEGRATIONS:

api/tests/unit/util/mappers/test_unit_mappers_sdk.py

@@ -36,8 +36,15 @@ def test_map_environment_to_sdk_document__return_expected(
     identity_featurestate.enabled = True
     identity_featurestate.save()
 
+    identities_with_overrides = (
+        Identity.objects.only_overrides().filter(environment=environment).all()
+    )
+
     # When
-    result = map_environment_to_sdk_document(environment)
+    result = map_environment_to_sdk_document(
+        environment,
+        identities_with_overrides=identities_with_overrides,
+    )
 
     # Then
     assert result == {

api/util/mappers/dynamodb.py

@@ -48,6 +48,7 @@ def map_environment_to_environment_document(
         field_name: _map_value_to_document_value(value)
         for field_name, value in map_environment_to_engine(
             environment,
+            with_integrations=True,
         )
     }
 

api/util/mappers/engine.py

@@ -49,6 +49,16 @@
     "map_traits_to_engine",
 )
 
+INTEGRATIONS_RELATION_NAMES = [
+    "amplitude_config",
+    "dynatrace_config",
+    "heap_config",
+    "mixpanel_config",
+    "rudderstack_config",
+    "segment_config",
+    "webhook_config",
+]
+
 
 def map_traits_to_engine(traits: Iterable["Trait"]) -> list[TraitModel]:
     return [
@@ -174,6 +184,8 @@ def map_mv_option_to_engine(
 
 def map_environment_to_engine(
     environment: "Environment",
+    *,
+    with_integrations: bool = True,
 ) -> EnvironmentModel:
     """
     Maps Core API's `environments.models.Environment` model instance to the
@@ -215,22 +227,14 @@ def map_environment_to_engine(
     }
 
     # Read integrations.
-    integration_configs: dict[str, Optional["EnvironmentIntegrationModel"]] = {}
-    for attr_name in (
-        "amplitude_config",
-        "dynatrace_config",
-        "heap_config",
-        "mixpanel_config",
-        "rudderstack_config",
-        "segment_config",
-    ):
-        integration_config = getattr(environment, attr_name, None)
-        if integration_config and not integration_config.deleted:
-            integration_configs[attr_name] = integration_config
-
-    webhook_config: Optional["WebhookConfiguration"] = getattr(
-        environment, "webhook_config", None
-    )
+    integration_configs: dict[
+        str, "EnvironmentIntegrationModel | WebhookConfiguration | None"
+    ] = {}
+    if with_integrations:
+        for attr_name in INTEGRATIONS_RELATION_NAMES:
+            integration_config = getattr(environment, attr_name, None)
+            if integration_config and not integration_config.deleted:
+                integration_configs[attr_name] = integration_config
 
     # No reading from ORM past this point!
 
@@ -306,13 +310,8 @@ def map_environment_to_engine(
     segment_config_model = map_integration_to_engine(
         integration_configs.pop("segment_config", None),
     )
-
-    webhook_config_model = (
-        map_webhook_config_to_engine(
-            webhook_config,
-        )
-        if webhook_config and not webhook_config.deleted
-        else None
+    webhook_config_model = map_webhook_config_to_engine(
+        integration_configs.pop("webhook_config", None),
     )
 
     return EnvironmentModel(

api/util/mappers/sdk.py

@@ -5,7 +5,8 @@
     map_identity_to_engine,
 )
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
+    from environments.identities.models import Identity
     from environments.models import Environment
 
 ENVIRONMENT_RESPONSE_EXCLUDE_FIELDS = [
@@ -25,6 +26,8 @@
 
 def map_environment_to_sdk_document(
     environment: "Environment",
+    *,
+    identities_with_overrides: list["Identity"] | None = None,
 ) -> SDKDocument:
     """
     Map an `environments.models.Environment` instance to an SDK document
@@ -33,17 +36,14 @@ def map_environment_to_sdk_document(
     It's virtually the same data that gets indexed in DynamoDB,
     except it presents identity overrides and omits integrations configurations.
     """
-    # Read relationships.
-    identities_with_overrides = environment.identities.with_overrides().all()
-
     # Get the engine data.
-    engine_environment = map_environment_to_engine(environment)
+    engine_environment = map_environment_to_engine(environment, with_integrations=False)
 
     # No reading from ORM past this point!
 
     # Prepare relationships.
     engine_environment.identity_overrides = [
-        map_identity_to_engine(identity) for identity in identities_with_overrides
+        map_identity_to_engine(identity) for identity in identities_with_overrides or []
     ]
 
     return engine_environment.model_dump(