feat: Identity overrides in environment document (#3766)

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/managers.py

@@ -6,33 +6,21 @@
 
 
 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(
-                    "feature_states",
-                    queryset=FeatureState.objects.select_related(
-                        "feature", "feature_state_value"
-                    ),
-                ),
-                Prefetch(
-                    "feature_states__multivariate_feature_state_values",
-                    queryset=MultivariateFeatureStateValue.objects.select_related(
-                        "multivariate_feature_option"
-                    ),
-                ),
                 "project__segments",
                 "project__segments__rules",
                 "project__segments__rules__rules",
@@ -55,6 +43,7 @@ def filter_for_document_builder(self, *args, **kwargs):
                         "multivariate_feature_option"
                     ),
                 ),
+                *extra_prefetch_related or (),
             )
             .filter(*args, **kwargs)
         )

api/environments/models.py

@@ -1,6 +1,3 @@
-# -*- coding: utf-8 -*-
-from __future__ import unicode_literals
-
 import logging
 import typing
 from copy import deepcopy
@@ -11,7 +8,7 @@
 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 (
@@ -35,6 +32,7 @@
     generate_client_api_key,
     generate_server_api_key,
 )
+from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES
 from environments.dynamodb import (
     DynamoEnvironmentAPIKeyWrapper,
     DynamoEnvironmentV2Wrapper,
@@ -43,10 +41,11 @@
 from environments.exceptions import EnvironmentHeaderNotPresentError
 from environments.managers import EnvironmentManager
 from features.models import Feature, FeatureSegment, FeatureState
+from features.multivariate.models import MultivariateFeatureStateValue
 from metadata.models import Metadata
 from projects.models import IdentityOverridesV2MigrationStatus, Project
 from segments.models import Segment
-from util.mappers import map_environment_to_environment_document
+from util.mappers import map_environment_to_sdk_document
 from webhooks.models import AbstractBaseExportableWebhookModel
 
 logger = logging.getLogger(__name__)
@@ -207,11 +206,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"
@@ -237,7 +232,24 @@ 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,
+                extra_prefetch_related=[
+                    Prefetch(
+                        "feature_states",
+                        queryset=FeatureState.objects.select_related(
+                            "feature", "feature_state_value"
+                        ),
+                    ),
+                    Prefetch(
+                        "feature_states__multivariate_feature_state_values",
+                        queryset=MultivariateFeatureStateValue.objects.select_related(
+                            "multivariate_feature_option"
+                        ),
+                    ),
+                ],
+            )
         )
         if not environments:
             return
@@ -363,8 +375,40 @@ 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_environment_document(environment)
+        environment = cls.objects.filter_for_document_builder(
+            api_key=api_key,
+            extra_prefetch_related=[
+                Prefetch(
+                    "feature_states",
+                    queryset=FeatureState.objects.select_related(
+                        "feature",
+                        "feature_state_value",
+                        "identity",
+                        "identity__environment",
+                    ).prefetch_related(
+                        Prefetch(
+                            "identity__identity_features",
+                            queryset=FeatureState.objects.select_related(
+                                "feature", "feature_state_value", "environment"
+                            ),
+                        ),
+                        Prefetch(
+                            "identity__identity_features__multivariate_feature_state_values",
+                            queryset=MultivariateFeatureStateValue.objects.select_related(
+                                "multivariate_feature_option"
+                            ),
+                        ),
+                    ),
+                ),
+                Prefetch(
+                    "feature_states__multivariate_feature_state_values",
+                    queryset=MultivariateFeatureStateValue.objects.select_related(
+                        "multivariate_feature_option"
+                    ),
+                ),
+            ],
+        ).get()
+        return map_environment_to_sdk_document(environment)
 
     def _get_environment(self):
         return self

api/environments/sdk/schemas.py

@@ -1,14 +1,10 @@
 from flag_engine.environments.models import EnvironmentModel
 
+from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES
 from util.pydantic import exclude_model_fields
 
 SDKEnvironmentDocumentModel = exclude_model_fields(
     EnvironmentModel,
-    "amplitude_config",
+    *IDENTITY_INTEGRATIONS_RELATION_NAMES,
     "dynatrace_config",
-    "heap_config",
-    "mixpanel_config",
-    "rudderstack_config",
-    "segment_config",
-    "webhook_config",
 )

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/environments/test_unit_environments_views_sdk_environment.py

@@ -1,19 +1,36 @@
+from typing import TYPE_CHECKING
+
 from core.constants import FLAGSMITH_UPDATED_AT_HEADER
 from django.urls import reverse
 from flag_engine.segments.constants import EQUAL
 from rest_framework import status
 from rest_framework.test import APIClient
 
+from environments.identities.models import Identity
 from environments.models import Environment, EnvironmentAPIKey
 from features.feature_types import MULTIVARIATE
-from features.models import Feature, FeatureSegment, FeatureState
+from features.models import (
+    STRING,
+    Feature,
+    FeatureSegment,
+    FeatureState,
+    FeatureStateValue,
+)
 from features.multivariate.models import MultivariateFeatureOption
 from segments.models import Condition, Segment, SegmentRule
 
+if TYPE_CHECKING:
+    from pytest_django import DjangoAssertNumQueries
+
+    from organisations.models import Organisation
+    from projects.models import Project
+
 
 def test_get_environment_document(
-    organisation_one, organisation_one_project_one, django_assert_num_queries
-):
+    organisation_one: "Organisation",
+    organisation_one_project_one: "Project",
+    django_assert_num_queries: "DjangoAssertNumQueries",
+) -> None:
     # Given
     project = organisation_one_project_one
 
@@ -57,6 +74,21 @@ def test_get_environment_document(
             enabled=True,
         )
 
+        # Add identity overrides
+        identity = Identity.objects.create(
+            environment=environment,
+            identifier=f"identity_{i}",
+        )
+        identity_feature_state = FeatureState.objects.create(
+            identity=identity,
+            feature=feature,
+            environment=environment,
+        )
+        FeatureStateValue.objects.filter(feature_state=identity_feature_state).update(
+            string_value="overridden",
+            type=STRING,
+        )
+
     for i in range(10):
         mv_feature = Feature.objects.create(
             name=f"mv_feature_{i}", project=project, type=MULTIVARIATE
@@ -76,7 +108,7 @@ def test_get_environment_document(
     url = reverse("api-v1:environment-document")
 
     # When
-    with django_assert_num_queries(13):
+    with django_assert_num_queries(15):
         response = client.get(url)
 
     # Then
@@ -88,8 +120,9 @@ def test_get_environment_document(
 
 
 def test_get_environment_document_fails_with_invalid_key(
-    organisation_one, organisation_one_project_one
-):
+    organisation_one: "Organisation",
+    organisation_one_project_one: "Project",
+) -> None:
     # Given
     project = organisation_one_project_one