From e8d133726c5b724057175bfc7d945e12a733f9d1 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Wed, 15 May 2024 15:57:00 +0100 Subject: [PATCH] feat: Identity overrides in environment document (#3766) --- api/environments/constants.py | 8 ++ api/environments/managers.py | 29 ++--- api/environments/models.py | 70 +++++++++--- api/environments/sdk/schemas.py | 8 +- api/integrations/integration.py | 16 ++- ...unit_environments_views_sdk_environment.py | 45 +++++++- .../util/mappers/test_unit_mappers_sdk.py | 108 ++++++++++++++++++ api/util/mappers/__init__.py | 2 + api/util/mappers/dynamodb.py | 1 + api/util/mappers/engine.py | 43 +++---- api/util/mappers/sdk.py | 51 +++++++++ 11 files changed, 306 insertions(+), 75 deletions(-) create mode 100644 api/environments/constants.py create mode 100644 api/tests/unit/util/mappers/test_unit_mappers_sdk.py create mode 100644 api/util/mappers/sdk.py diff --git a/api/environments/constants.py b/api/environments/constants.py new file mode 100644 index 000000000000..8f83d38c5b3e --- /dev/null +++ b/api/environments/constants.py @@ -0,0 +1,8 @@ +IDENTITY_INTEGRATIONS_RELATION_NAMES = [ + "amplitude_config", + "heap_config", + "mixpanel_config", + "rudderstack_config", + "segment_config", + "webhook_config", +] diff --git a/api/environments/managers.py b/api/environments/managers.py index bf9d7cf03338..fe0e0a4b51df 100644 --- a/api/environments/managers.py +++ b/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) ) diff --git a/api/environments/models.py b/api/environments/models.py index 2cec24c57594..235128c870f9 100644 --- a/api/environments/models.py +++ b/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 diff --git a/api/environments/sdk/schemas.py b/api/environments/sdk/schemas.py index 9741f4130df2..f96b656fe81d 100644 --- a/api/environments/sdk/schemas.py +++ b/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", ) diff --git a/api/integrations/integration.py b/api/integrations/integration.py index faecf80aebff..f6702f86af4b 100644 --- a/api/integrations/integration.py +++ b/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: diff --git a/api/tests/unit/environments/test_unit_environments_views_sdk_environment.py b/api/tests/unit/environments/test_unit_environments_views_sdk_environment.py index 70192c304ec2..c14b976d5bcd 100644 --- a/api/tests/unit/environments/test_unit_environments_views_sdk_environment.py +++ b/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 diff --git a/api/tests/unit/util/mappers/test_unit_mappers_sdk.py b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py new file mode 100644 index 000000000000..7363f1e840db --- /dev/null +++ b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py @@ -0,0 +1,108 @@ +from typing import TYPE_CHECKING + +import pytest + +from environments.identities.models import Identity +from util.mappers.sdk import map_environment_to_sdk_document + +if TYPE_CHECKING: # pragma: no cover + from pytest_mock import MockerFixture + + from environments.models import Environment + from features.models import FeatureState + + +@pytest.fixture() +def identity_without_overrides(environment): + return Identity.objects.create( + identifier="test_identity_without_overrides", environment=environment + ) + + +def test_map_environment_to_sdk_document__return_expected( + mocker: "MockerFixture", + environment: "Environment", + feature_state: "FeatureState", + identity: Identity, + identity_featurestate: "FeatureState", + identity_without_overrides: Identity, +) -> None: + # Given + expected_overridden_value = "some-overridden-value" + expected_api_key = environment.api_key + + identity_featurestate.feature_state_value.string_value = expected_overridden_value + identity_featurestate.feature_state_value.save() + identity_featurestate.enabled = True + identity_featurestate.save() + + # When + result = map_environment_to_sdk_document(environment) + + # Then + assert result == { + "allow_client_traits": True, + "api_key": expected_api_key, + "feature_states": [ + { + "django_id": feature_state.pk, + "enabled": False, + "feature": { + "id": feature_state.feature.pk, + "name": "Test Feature1", + "type": "STANDARD", + }, + "feature_segment": None, + "feature_state_value": None, + "featurestate_uuid": feature_state.uuid, + "multivariate_feature_state_values": [], + } + ], + "identity_overrides": [ + { + "composite_key": identity.composite_key, + "created_date": identity.created_date, + "django_id": identity.pk, + "environment_api_key": expected_api_key, + "identifier": identity.identifier, + "identity_features": [ + { + "django_id": identity_featurestate.pk, + "enabled": True, + "feature": { + "id": identity_featurestate.feature.pk, + "name": identity_featurestate.feature.name, + "type": identity_featurestate.feature.type, + }, + "feature_segment": None, + "feature_state_value": expected_overridden_value, + "featurestate_uuid": identity_featurestate.uuid, + "multivariate_feature_state_values": [], + } + ], + "identity_traits": [], + "identity_uuid": mocker.ANY, + } + ], + "hide_disabled_flags": None, + "hide_sensitive_data": False, + "id": environment.pk, + "name": "Test Environment", + "project": { + "enable_realtime_updates": False, + "hide_disabled_flags": False, + "id": environment.project.pk, + "name": "Test Project", + "organisation": { + "feature_analytics": False, + "id": environment.project.organisation.pk, + "name": "Test Org", + "persist_trait_data": True, + "stop_serving_flags": False, + }, + "segments": [], + "server_key_only_feature_ids": [], + }, + "updated_at": environment.updated_at, + "use_identity_composite_key_for_hashing": True, + } diff --git a/api/util/mappers/__init__.py b/api/util/mappers/__init__.py index 995f02308be2..6b5f3d2a7f41 100644 --- a/api/util/mappers/__init__.py +++ b/api/util/mappers/__init__.py @@ -14,6 +14,7 @@ map_identity_to_engine, map_mv_option_to_engine, ) +from util.mappers.sdk import map_environment_to_sdk_document __all__ = ( "map_engine_feature_state_to_identity_override", @@ -21,6 +22,7 @@ "map_environment_api_key_to_environment_api_key_document", "map_environment_to_environment_document", "map_environment_to_environment_v2_document", + "map_environment_to_sdk_document", "map_feature_state_to_engine", "map_feature_to_engine", "map_identity_changeset_to_identity_override_changeset", diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index 95a0850f9d70..b4da6bcdb903 100644 --- a/api/util/mappers/dynamodb.py +++ b/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, ) } diff --git a/api/util/mappers/engine.py b/api/util/mappers/engine.py index 5730dc9846c2..243b1fff6eec 100644 --- a/api/util/mappers/engine.py +++ b/api/util/mappers/engine.py @@ -24,7 +24,9 @@ SegmentRuleModel, ) -if TYPE_CHECKING: +from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES + +if TYPE_CHECKING: # pragma: no cover from environments.identities.models import Identity, Trait from environments.models import Environment, EnvironmentAPIKey from features.models import Feature, FeatureSegment, FeatureState @@ -174,6 +176,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 +219,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 IDENTITY_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! @@ -291,9 +287,6 @@ def map_environment_to_engine( amplitude_config_model = map_integration_to_engine( integration_configs.pop("amplitude_config", None), ) - dynatrace_config_model = map_integration_to_engine( - integration_configs.pop("dynatrace_config", None), - ) heap_config_model = map_integration_to_engine( integration_configs.pop("heap_config", None), ) @@ -306,13 +299,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( @@ -333,7 +321,6 @@ def map_environment_to_engine( # # Integrations: amplitude_config=amplitude_config_model, - dynatrace_config=dynatrace_config_model, heap_config=heap_config_model, mixpanel_config=mixpanel_config_model, rudderstack_config=rudderstack_config_model, diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py new file mode 100644 index 000000000000..bc94def7ae8f --- /dev/null +++ b/api/util/mappers/sdk.py @@ -0,0 +1,51 @@ +from typing import TYPE_CHECKING, TypeAlias + +from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES +from util.mappers.engine import ( + map_environment_to_engine, + map_identity_to_engine, +) + +if TYPE_CHECKING: # pragma: no cover + from environments.models import Environment + + +SDKDocumentValue: TypeAlias = dict[str, "SDKDocumentValue"] | str | bool | None | float +SDKDocument: TypeAlias = dict[str, SDKDocumentValue] + +SDK_DOCUMENT_EXCLUDE = [ + *IDENTITY_INTEGRATIONS_RELATION_NAMES, + "dynatrace_config", +] + + +def map_environment_to_sdk_document(environment: "Environment") -> SDKDocument: + """ + Map an `environments.models.Environment` instance to an SDK document + used by SDKs with local evaluation mode. + + 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 = {} + for feature_state in environment.feature_states.all(): + if (identity_id := feature_state.identity_id) and ( + identity_id not in identities_with_overrides + ): + identities_with_overrides[identity_id] = feature_state.identity + + # Get the engine data. + 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, with_traits=False) + for identity in identities_with_overrides.values() + ] + + return engine_environment.model_dump( + exclude=SDK_DOCUMENT_EXCLUDE, + )