From aa1ed78d480eaf310986713d3e71b7a439ad95d6 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Fri, 25 Sep 2026 21:37:38 +0200 Subject: [PATCH 1/5] refactor(api): map segments for membership counts with the evaluation mapper --- api/segment_membership/services.py | 7 +- api/util/engine_models/context/__init__.py | 0 api/util/engine_models/context/mappers.py | 116 --------------------- 3 files changed, 3 insertions(+), 120 deletions(-) delete mode 100644 api/util/engine_models/context/__init__.py delete mode 100644 api/util/engine_models/context/mappers.py diff --git a/api/segment_membership/services.py b/api/segment_membership/services.py index c68295fc088b..24cb9f2640ac 100644 --- a/api/segment_membership/services.py +++ b/api/segment_membership/services.py @@ -17,14 +17,13 @@ from task_processor.models import Task from environments.models import Environment +from evaluation.mappers import map_segment_to_segment_context from integrations.flagsmith.client import get_openfeature_client from organisations.models import Organisation from projects.models import Project from segment_membership.models import SegmentMembershipCount from segment_membership.types import ClickHouseReadIdentityRow, SegmentMember from segments.models import Segment -from util.engine_models.context.mappers import map_segment_to_segment_context -from util.mappers.engine import map_segment_to_engine logger = structlog.get_logger("segment_membership") @@ -142,7 +141,7 @@ def compute_segment_counts_for_project( binder=binder, ) predicate = translate_segment( - map_segment_to_segment_context(map_segment_to_engine(seg)), + map_segment_to_segment_context(seg), # type: ignore[arg-type] translate_ctx, ) if predicate is None: @@ -208,7 +207,7 @@ def get_segment_members_page( binder=binder, ) predicate = translate_segment( - map_segment_to_segment_context(map_segment_to_engine(segment)), + map_segment_to_segment_context(segment), # type: ignore[arg-type] translate_ctx, ) if predicate is None: diff --git a/api/util/engine_models/context/__init__.py b/api/util/engine_models/context/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/context/mappers.py b/api/util/engine_models/context/mappers.py deleted file mode 100644 index 86a243d5774e..000000000000 --- a/api/util/engine_models/context/mappers.py +++ /dev/null @@ -1,116 +0,0 @@ -""" -Vendored and adapted mappers from flagsmith-flag-engine's fix/missing-export branch. - -The original `map_environment_identity_to_context` function has been adapted to -return v10's EvaluationContext TypedDict instead of the original return type. -""" - -import typing - -from flag_engine.context.types import ( - FeatureContext, - SegmentContext, - SegmentRule, -) - -from util.engine_models.features.models import ( - FeatureStateModel, - MultivariateFeatureStateValueModel, -) -from util.engine_models.segments.models import SegmentModel, SegmentRuleModel - - -def _map_feature_states_to_feature_contexts( - feature_states: typing.List[FeatureStateModel], -) -> typing.Dict[str, FeatureContext]: - """ - Map feature states to feature contexts. - - :param feature_states: A list of FeatureStateModel objects. - :return: A dictionary mapping feature names to their contexts. - """ - features: typing.Dict[str, FeatureContext] = {} - for feature_state in feature_states: - feature_context: FeatureContext = { - "key": str(feature_state.django_id or feature_state.featurestate_uuid), - "name": feature_state.feature.name, - "enabled": feature_state.enabled, - "value": feature_state.feature_state_value, - } - multivariate_feature_state_values: typing.List[ - MultivariateFeatureStateValueModel - ] - if multivariate_feature_state_values := list( - feature_state.multivariate_feature_state_values - ): - sorted_mv_values = sorted( - multivariate_feature_state_values, - key=_get_multivariate_feature_state_value_id, - ) - feature_context["variants"] = [ - { - "value": mv_value.multivariate_feature_option.value, - "weight": mv_value.percentage_allocation, - "priority": idx, - } - for idx, mv_value in enumerate(sorted_mv_values) - ] - if feature_segment := feature_state.feature_segment: - if (priority := feature_segment.priority) is not None: - feature_context["priority"] = priority - features[feature_state.feature.name] = feature_context - return features - - -def _map_segment_rules_to_segment_context_rules( - rules: typing.List[SegmentRuleModel], -) -> typing.List[SegmentRule]: - """ - Map segment rules to segment rules for the evaluation context. - - :param rules: A list of SegmentRuleModel objects. - :return: A list of SegmentRule objects. - """ - return [ - { - "type": rule.type, - "conditions": [ - { - "property": condition.property_ or "", - "operator": condition.operator, - "value": condition.value or "", - } - for condition in rule.conditions - ], - "rules": _map_segment_rules_to_segment_context_rules(rule.rules), - } - for rule in rules - ] - - -def _get_multivariate_feature_state_value_id( - multivariate_feature_state_value: MultivariateFeatureStateValueModel, -) -> int: - return ( - multivariate_feature_state_value.id - or multivariate_feature_state_value.mv_fs_value_uuid.int - ) - - -def map_segment_to_segment_context(segment: SegmentModel) -> SegmentContext: - """ - Map a SegmentModel Pydantic model to a SegmentContext TypedDict. - - :param segment: The SegmentModel object. - :return: A SegmentContext TypedDict. - """ - segment_ctx: SegmentContext = { - "key": str(segment.id), - "name": segment.name, - "rules": _map_segment_rules_to_segment_context_rules(segment.rules), - } - if segment_feature_states := segment.feature_states: - segment_ctx["overrides"] = list( - _map_feature_states_to_feature_contexts(segment_feature_states).values() - ) - return segment_ctx From 04ac5f8af0670b616a154ebd6ab58dc34639af56 Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Fri, 25 Sep 2026 22:02:32 +0200 Subject: [PATCH 2/5] refactor(api): build environment documents through flagsmith_schemas type adapters The environment and API key mappers now return plain dicts. Dynamo documents are validated through the flagsmith_schemas TypedDicts, which coerce their values, rather than through vendored models and a custom encoder. The uncompressed environment document no longer carries the always-empty `identity_overrides`, or `entity_selector` on non-Dynatrace integrations, as the compressed one already didn't. --- .../test_unit_features_features_service.py | 17 +- .../mappers/test_unit_mappers_dynamodb.py | 11 +- .../util/mappers/test_unit_mappers_engine.py | 370 ++++++++---------- api/util/mappers/dynamodb.py | 40 +- api/util/mappers/engine.py | 273 ++++++------- api/util/mappers/sdk.py | 23 +- 6 files changed, 348 insertions(+), 386 deletions(-) diff --git a/api/tests/unit/features/test_unit_features_features_service.py b/api/tests/unit/features/test_unit_features_features_service.py index fe1154a3a6fa..a046f2936717 100644 --- a/api/tests/unit/features/test_unit_features_features_service.py +++ b/api/tests/unit/features/test_unit_features_features_service.py @@ -13,6 +13,7 @@ from features.models import Feature, FeatureSegment, FeatureState from projects.models import EdgeV2MigrationStatus from users.models import FFAdminUser +from util.engine_models.features.models import FeatureStateModel from util.mappers.engine import ( map_feature_state_to_engine, map_identity_to_engine, @@ -247,10 +248,14 @@ def test_get_edge_overrides_data__multiple_overrides__returns_correct_counts( # replicate identity to Edge edge_identity = EdgeIdentity(map_identity_to_engine(identity, with_overrides=False)) edge_identity.add_feature_override( - map_feature_state_to_engine(identity_featurestate), + FeatureStateModel.model_validate( + map_feature_state_to_engine(identity_featurestate) + ), ) edge_identity.add_feature_override( - map_feature_state_to_engine(distinct_identity_featurestate), + FeatureStateModel.model_validate( + map_feature_state_to_engine(distinct_identity_featurestate) + ), ) edge_identity.save(admin_user) @@ -303,10 +308,14 @@ def test_get_edge_overrides_data__deleted_feature__skips_deleted( # type: ignor edge_identity = EdgeIdentity(map_identity_to_engine(identity, with_overrides=False)) # Create identity override for two different features edge_identity.add_feature_override( - map_feature_state_to_engine(identity_featurestate), + FeatureStateModel.model_validate( + map_feature_state_to_engine(identity_featurestate) + ), ) edge_identity.add_feature_override( - map_feature_state_to_engine(distinct_identity_featurestate), + FeatureStateModel.model_validate( + map_feature_state_to_engine(distinct_identity_featurestate) + ), ) edge_identity.save(admin_user) diff --git a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py index d96519bdd048..9351ec3ce76d 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py @@ -9,6 +9,7 @@ from environments.dynamodb.constants import ( ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, ) +from util.engine_models.features.models import FeatureStateModel from util.engine_models.identities.models import IdentityModel from util.mappers import dynamodb from util.mappers.engine import map_feature_state_to_engine @@ -56,7 +57,6 @@ def test_map_environment_to_environment_document__valid_environment__returns_exp "multivariate_feature_state_values": [], } ], - "identity_overrides": [], "heap_config": None, "hide_disabled_flags": None, "hide_sensitive_data": False, @@ -228,7 +228,6 @@ def test_map_environment_to_environment_v2_document__valid_environment__returns_ "allow_client_traits": True, "amplitude_config": None, "dynatrace_config": None, - "identity_overrides": [], "feature_states": [ { "django_id": Decimal(feature_state.pk), @@ -282,8 +281,12 @@ def test_map_identity_override_to_identity_override_document__decimal_feature_st # Given expected_feature_state_value = Decimal("1.111") - engine_feature_state = map_feature_state_to_engine(identity_featurestate) - engine_feature_state.feature_state_value = expected_feature_state_value + engine_feature_state = FeatureStateModel.model_validate( + { + **map_feature_state_to_engine(identity_featurestate), + "feature_state_value": expected_feature_state_value, + } + ) identity_override = dynamodb.map_engine_feature_state_to_identity_override( feature_state=engine_feature_state, identity_uuid=str(uuid.uuid4()), diff --git a/api/tests/unit/util/mappers/test_unit_mappers_engine.py b/api/tests/unit/util/mappers/test_unit_mappers_engine.py index 5f1f0a96d0ed..67af2beb8976 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_engine.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_engine.py @@ -20,31 +20,16 @@ from integrations.webhook.models import WebhookConfiguration from segments.models import Segment, SegmentRule from users.models import FFAdminUser -from util.engine_models.environments.integrations.models import IntegrationModel -from util.engine_models.environments.models import ( - EnvironmentAPIKeyModel, - EnvironmentModel, - WebhookModel, -) from util.engine_models.features.models import ( FeatureModel, - FeatureSegmentModel, FeatureStateModel, MultivariateFeatureOptionModel, - MultivariateFeatureStateValueModel, ) from util.engine_models.identities.models import ( IdentityFeaturesList, IdentityModel, ) from util.engine_models.identities.traits.models import TraitModel -from util.engine_models.organisations.models import OrganisationModel -from util.engine_models.projects.models import ProjectModel -from util.engine_models.segments.models import ( - SegmentConditionModel, - SegmentModel, - SegmentRuleModel, -) from util.mappers import engine if TYPE_CHECKING: @@ -113,23 +98,11 @@ def test_map_segment_rule_to_engine__nested_rule__returns_expected_model( result = engine.map_segment_rule_to_engine(matching_rule) # Then - assert result == SegmentRuleModel( - type="ALL", - rules=[ - SegmentRuleModel( - type="ALL", - rules=[], - conditions=[], - ) - ], - conditions=[ - SegmentConditionModel( - operator="EQUAL", - value="value1", - property_="key1", - ) - ], - ) + assert result == { + "type": "ALL", + "rules": [{"type": "ALL", "rules": [], "conditions": []}], + "conditions": [{"operator": "EQUAL", "value": "value1", "property_": "key1"}], + } def test_map_integration_to_engine__valid_integration__returns_expected_model() -> None: @@ -142,7 +115,11 @@ class Meta: api_key = "test" integration = TestIntegration(base_url=base_url, api_key=api_key) - expected_result = IntegrationModel(base_url=base_url, api_key=api_key) + expected_result = { + "base_url": base_url, + "api_key": api_key, + "entity_selector": None, + } # When result = engine.map_integration_to_engine(integration) @@ -170,11 +147,11 @@ def test_map_integration_to_engine__dynatrace__return_expected() -> None: api_key=api_key, entity_selector=entity_selector, ) - expected_result = IntegrationModel( - base_url=base_url, - api_key=api_key, - entity_selector=entity_selector, - ) + expected_result = { + "base_url": base_url, + "api_key": api_key, + "entity_selector": entity_selector, + } # When result = engine.map_integration_to_engine(integration) @@ -192,10 +169,7 @@ def test_map_webhook_config_to_engine__valid_config__returns_expected_model() -> url=url, secret=secret, ) - expected_result = WebhookModel( - url=url, - secret=secret, - ) + expected_result = {"url": url, "secret": secret} # When result = engine.map_webhook_config_to_engine(webhook_config) @@ -217,19 +191,15 @@ def test_map_feature_state_to_engine__standard_feature__returns_expected_model( feature_state: FeatureState, ) -> None: # Given - expected_result = FeatureStateModel( - feature=FeatureModel( - id=feature.id, - name="Test Feature1", - type="STANDARD", - ), - enabled=False, - django_id=feature_state.id, - feature_segment=None, - featurestate_uuid=feature_state.uuid, - feature_state_value=None, - multivariate_feature_state_values=[], # type: ignore[arg-type] - ) + expected_result = { + "feature": {"id": feature.id, "name": "Test Feature1", "type": "STANDARD"}, + "enabled": False, + "django_id": feature_state.id, + "feature_segment": None, + "featurestate_uuid": feature_state.uuid, + "feature_state_value": None, + "multivariate_feature_state_values": [], + } # When result = engine.map_feature_state_to_engine( @@ -252,7 +222,7 @@ def test_map_feature_state_to_engine__mv_hashing_salt_set__uses_salt_as_django_i # Then the salt is used as the engine document's django_id so that variant # bucketing stays stable across feature state recreation - assert result.django_id == feature_state.mv_hashing_salt + assert result["django_id"] == feature_state.mv_hashing_salt def test_map_feature_state_to_engine__feature_segment__return_expected( @@ -263,31 +233,32 @@ def test_map_feature_state_to_engine__feature_segment__return_expected( mv_fs_value = ( segment_multivariate_feature_state.multivariate_feature_state_values.get() ) - expected_result = FeatureStateModel( - feature=FeatureModel( - id=multivariate_feature.id, - name="feature", - type="MULTIVARIATE", - ), - enabled=False, - django_id=segment_multivariate_feature_state.id, - feature_segment=FeatureSegmentModel( - priority=segment_multivariate_feature_state.feature_segment.priority, # type: ignore[union-attr] - ), - featurestate_uuid=segment_multivariate_feature_state.uuid, - feature_state_value="control", - multivariate_feature_state_values=[ # type: ignore[arg-type] - MultivariateFeatureStateValueModel( - multivariate_feature_option=MultivariateFeatureOptionModel( - value=mv_fs_value.multivariate_feature_option.value, - id=mv_fs_value.multivariate_feature_option.id, - ), - percentage_allocation=mv_fs_value.percentage_allocation, - id=mv_fs_value.id, - mv_fs_value_uuid=mv_fs_value.uuid, - ), + expected_result = { + "feature": { + "id": multivariate_feature.id, + "name": "feature", + "type": "MULTIVARIATE", + }, + "enabled": False, + "django_id": segment_multivariate_feature_state.id, + "feature_segment": { + "priority": segment_multivariate_feature_state.feature_segment.priority, # type: ignore[union-attr] + }, + "featurestate_uuid": segment_multivariate_feature_state.uuid, + "feature_state_value": "control", + "multivariate_feature_state_values": [ + { + "multivariate_feature_option": { + "value": mv_fs_value.multivariate_feature_option.value, + "id": mv_fs_value.multivariate_feature_option.id, + "key": None, + }, + "percentage_allocation": mv_fs_value.percentage_allocation, + "id": mv_fs_value.id, + "mv_fs_value_uuid": mv_fs_value.uuid, + }, ], - ) + } # When result = engine.map_feature_state_to_engine( @@ -312,12 +283,7 @@ def test_map_mv_option_to_engine__option_with_key__includes_key( result = engine.map_mv_option_to_engine(mv_option) # Then - assert result.key == "control" - assert result == MultivariateFeatureOptionModel( - value=mv_option.value, - id=mv_option.id, - key="control", - ) + assert result == {"value": mv_option.value, "id": mv_option.id, "key": "control"} def test_map_feature_state_to_engine__mv_option_with_key__key_in_serialised_document( @@ -338,11 +304,10 @@ def test_map_feature_state_to_engine__mv_option_with_key__key_in_serialised_docu segment_multivariate_feature_state, mv_fs_values=[mv_fs_value], ) - document = result.model_dump() # Then assert ( - document["multivariate_feature_state_values"][0]["multivariate_feature_option"][ + result["multivariate_feature_state_values"][0]["multivariate_feature_option"][ "key" ] == "control" @@ -411,88 +376,92 @@ def test_map_environment_to_engine__multiple_segments_and_versions__returns_expe ) deleted_segment_configuration.delete() - expected_feature_model = FeatureModel( - id=feature.id, name="Test Feature1", type="STANDARD" - ) - expected_segment_feature_state_model = FeatureStateModel( - feature=expected_feature_model, - enabled=False, - django_id=versioned_segment_feature_state.id, - feature_segment=FeatureSegmentModel( - priority=feature_segment.priority, - ), - featurestate_uuid=versioned_segment_feature_state.uuid, - feature_state_value=None, - multivariate_feature_state_values=[], # type: ignore[arg-type] - ) - expected_feature_state_model = FeatureStateModel( - feature=expected_feature_model, - enabled=False, - django_id=feature_state.id, - featurestate_uuid=feature_state.uuid, - feature_state_value=None, - multivariate_feature_state_values=[], # type: ignore[arg-type] - ) - expected_project_model = ProjectModel( - id=environment.project.id, - name="Test Project", - organisation=OrganisationModel( - id=environment.project.organisation.id, - name="Test Org", - feature_analytics=False, - stop_serving_flags=False, - persist_trait_data=True, - ), - hide_disabled_flags=False, - segments=[ - SegmentModel( - id=segment.id, - name=segment.name, - rules=[], - feature_states=[expected_segment_feature_state_model], - ), + expected_feature_model = { + "id": feature.id, + "name": "Test Feature1", + "type": "STANDARD", + } + expected_segment_feature_state_model = { + "feature": expected_feature_model, + "enabled": False, + "django_id": versioned_segment_feature_state.id, + "feature_segment": {"priority": feature_segment.priority}, + "featurestate_uuid": versioned_segment_feature_state.uuid, + "feature_state_value": None, + "multivariate_feature_state_values": [], + } + expected_feature_state_model = { + "feature": expected_feature_model, + "enabled": False, + "django_id": feature_state.id, + "feature_segment": None, + "featurestate_uuid": feature_state.uuid, + "feature_state_value": None, + "multivariate_feature_state_values": [], + } + expected_project_model = { + "id": environment.project.id, + "name": "Test Project", + "organisation": { + "id": environment.project.organisation.id, + "name": "Test Org", + "feature_analytics": False, + "stop_serving_flags": False, + "persist_trait_data": True, + }, + "hide_disabled_flags": False, + "segments": [ + { + "id": segment.id, + "name": segment.name, + "rules": [], + "feature_states": [expected_segment_feature_state_model], + }, ], - enable_realtime_updates=False, - server_key_only_feature_ids=[], - ) - - expected_result = EnvironmentModel( - id=environment.id, - api_key=environment.api_key, - project=expected_project_model, - feature_states=[expected_feature_state_model], - name=environment.name, - allow_client_traits=environment.allow_client_traits, - updated_at=environment.updated_at, - use_identity_composite_key_for_hashing=environment.use_identity_composite_key_for_hashing, - use_identity_overrides_in_local_eval=environment.use_identity_overrides_in_local_eval, - hide_sensitive_data=environment.hide_sensitive_data, - hide_disabled_flags=environment.hide_disabled_flags, - onboarding_pending=True, - amplitude_config=None, - dynatrace_config=None, - heap_config=None, - mixpanel_config={ # type: ignore[arg-type] + "enable_realtime_updates": False, + "server_key_only_feature_ids": [], + } + + expected_result = { + "id": environment.id, + "api_key": environment.api_key, + "project": expected_project_model, + "feature_states": [expected_feature_state_model], + "identity_overrides": [], + "name": environment.name, + "allow_client_traits": environment.allow_client_traits, + "updated_at": environment.updated_at, + "use_identity_composite_key_for_hashing": environment.use_identity_composite_key_for_hashing, + "use_identity_overrides_in_local_eval": environment.use_identity_overrides_in_local_eval, + "hide_sensitive_data": environment.hide_sensitive_data, + "hide_disabled_flags": environment.hide_disabled_flags, + "onboarding_pending": True, + "amplitude_config": None, + "dynatrace_config": None, + "heap_config": None, + "mixpanel_config": { "base_url": mixpanel_configuration.base_url, "api_key": mixpanel_configuration.api_key, + "entity_selector": None, }, - rudderstack_config=None, - segment_config=None, # note: segment configuration should not appear as it was deleted - webhook_config={ # type: ignore[arg-type] + "rudderstack_config": None, + "segment_config": None, # note: segment configuration should not appear as it was deleted + "webhook_config": { "url": webhook_configuration.url, "secret": webhook_configuration.secret, }, - ) + } # When result = engine.map_environment_to_engine(environment) segment_feature_state_uuids = [ - fs.featurestate_uuid for fs in result.project.segments[0].feature_states + fs["featurestate_uuid"] + for fs in result["project"]["segments"][0]["feature_states"] ] # Then - assert len(result.feature_states) == 1 - assert result.feature_states[0].django_id == feature_state.id + assert len(result["feature_states"]) == 1 + assert result["feature_states"][0]["django_id"] == feature_state.id assert result == expected_result @@ -517,7 +486,7 @@ def test_map_environment_to_engine__feature_specific_segment_not_in_env__exclude result = engine.map_environment_to_engine(environment) # Then - segment_ids = [s.id for s in result.project.segments] + segment_ids = [s["id"] for s in result["project"]["segments"]] assert feature_specific_segment.id not in segment_ids @@ -542,7 +511,7 @@ def test_map_environment_to_engine__feature_specific_segment_in_env__includes_se result = engine.map_environment_to_engine(environment) # Then - segment_ids = [s.id for s in result.project.segments] + segment_ids = [s["id"] for s in result["project"]["segments"]] assert feature_specific_segment.id in segment_ids @@ -559,7 +528,7 @@ def test_map_environment_to_engine__project_wide_segment_not_in_env__includes_se result = engine.map_environment_to_engine(environment) # Then - segment_ids = [s.id for s in result.project.segments] + segment_ids = [s["id"] for s in result["project"]["segments"]] assert segment.id in segment_ids @@ -574,15 +543,15 @@ def test_map_environment_api_key_to_engine__valid_key__returns_expected_model( result = engine.map_environment_api_key_to_engine(environment_api_key) # Then - assert result == EnvironmentAPIKeyModel( - id=environment_api_key.pk, - key=environment_api_key.key, - created_at=environment_api_key.created_at, - name=environment_api_key.name, - client_api_key=client_api_key, - expires_at=environment_api_key.expires_at, - active=environment_api_key.active, - ) + assert result == { + "id": environment_api_key.pk, + "key": environment_api_key.key, + "created_at": environment_api_key.created_at, + "name": environment_api_key.name, + "client_api_key": client_api_key, + "expires_at": environment_api_key.expires_at, + "active": environment_api_key.active, + } def test_map_identity_to_engine__identity_with_traits_and_overrides__returns_expected_model( @@ -675,8 +644,8 @@ def test_map_environment_to_engine__different_versions__returns_latest_live_from result = engine.map_environment_to_engine(environment) # Then - assert len(result.feature_states) == 1 - assert result.feature_states[0].django_id == v15_feature_state.id + assert len(result["feature_states"]) == 1 + assert result["feature_states"][0]["django_id"] == v15_feature_state.id def test_map_environment_to_engine__after_v2_versioning_migration__returns_latest_versions( @@ -723,26 +692,26 @@ def test_map_environment_to_engine__after_v2_versioning_migration__returns_lates # Then assert result - assert len(result.feature_states) == 1 - mapped_environment_feature_state = result.feature_states[0] + assert len(result["feature_states"]) == 1 + mapped_environment_feature_state = result["feature_states"][0] assert ( - mapped_environment_feature_state.featurestate_uuid + mapped_environment_feature_state["featurestate_uuid"] == v2_environment_feature_state.uuid ) - assert mapped_environment_feature_state.enabled is True + assert mapped_environment_feature_state["enabled"] is True assert ( - mapped_environment_feature_state.feature_state_value + mapped_environment_feature_state["feature_state_value"] == v2_environment_feature_state_value ) - assert len(result.project.segments) == 1 - assert len(result.project.segments[0].feature_states) == 1 + assert len(result["project"]["segments"]) == 1 + assert len(result["project"]["segments"][0]["feature_states"]) == 1 - mapped_segment_override = result.project.segments[0].feature_states[0] - assert mapped_segment_override.featurestate_uuid == v2_segment_override.uuid - assert mapped_segment_override.enabled is True - assert mapped_segment_override.feature_state_value == v2_segment_override_value + mapped_segment_override = result["project"]["segments"][0]["feature_states"][0] + assert mapped_segment_override["featurestate_uuid"] == v2_segment_override.uuid + assert mapped_segment_override["enabled"] is True + assert mapped_segment_override["feature_state_value"] == v2_segment_override_value def test_map_environment_to_engine__v2_versioning_segment_override_removed__returns_remaining_override( @@ -794,11 +763,11 @@ def test_map_environment_to_engine__v2_versioning_segment_override_removed__retu environment_model = engine.map_environment_to_engine(environment_v2_versioning) # Then - assert len(environment_model.project.segments[0].feature_states) == 1 - assert ( - environment_model.project.segments[0].feature_states[0].featurestate_uuid - == v3_segment_override.uuid - ) + segment_feature_states = environment_model["project"]["segments"][0][ + "feature_states" + ] + assert len(segment_feature_states) == 1 + assert segment_feature_states[0]["featurestate_uuid"] == v3_segment_override.uuid def test_map_environment_to_engine__running_experiment__stamps_every_state_of_feature( @@ -827,15 +796,15 @@ def test_map_environment_to_engine__running_experiment__stamps_every_state_of_fe # Then (default_state,) = [ - fs for fs in result.feature_states if fs.feature.id == feature.id + fs for fs in result["feature_states"] if fs["feature"]["id"] == feature.id ] - assert default_state.metadata == { + assert default_state["metadata"] == { "experiment": {**expected_experiment, "in_experiment": False}, } assert { - segment.name: fs.metadata - for segment in result.project.segments - for fs in segment.feature_states + segment["name"]: fs["metadata"] + for segment in result["project"]["segments"] + for fs in segment["feature_states"] } == { "Experiment rollout": { "experiment": {**expected_experiment, "in_experiment": True}, @@ -861,10 +830,9 @@ def test_map_environment_to_engine__running_experiment__other_features_unstamped # Then (other_state,) = [ - fs for fs in result.feature_states if fs.feature.id == other_feature.id + fs for fs in result["feature_states"] if fs["feature"]["id"] == other_feature.id ] - assert other_state.metadata is None - assert "metadata" not in other_state.dict() + assert "metadata" not in other_state @pytest.mark.parametrize( @@ -887,11 +855,11 @@ def test_map_environment_to_engine__experiment_not_running__no_metadata( result = engine.map_environment_to_engine(environment) # Then - assert all(fs.metadata is None for fs in result.feature_states) + assert all("metadata" not in fs for fs in result["feature_states"]) assert all( - fs.metadata is None - for segment in result.project.segments - for fs in segment.feature_states + "metadata" not in fs + for segment in result["project"]["segments"] + for fs in segment["feature_states"] ) @@ -908,7 +876,7 @@ def test_map_environment_to_engine__experiment_without_rollout_segment__no_enrol # Then assert not any( - fs.metadata["experiment"]["in_experiment"] # type: ignore[index] - for segment in result.project.segments - for fs in segment.feature_states + fs["metadata"]["experiment"]["in_experiment"] + for segment in result["project"]["segments"] + for fs in segment["feature_states"] ) diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index f59458772629..7b55ea6f1464 100644 --- a/api/util/mappers/dynamodb.py +++ b/api/util/mappers/dynamodb.py @@ -3,6 +3,8 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, TypeVar, Union, cast from flagsmith_schemas.dynamodb import ( + Environment, + EnvironmentAPIKey, EnvironmentCompressed, EnvironmentV2MetaCompressed, ) @@ -31,7 +33,8 @@ if TYPE_CHECKING: from environments.identities.models import Identity - from environments.models import Environment, EnvironmentAPIKey + from environments.models import Environment as EnvironmentModel + from environments.models import EnvironmentAPIKey as EnvironmentAPIKeyModel from util.engine_models.identities.models import IdentityModel @@ -46,6 +49,10 @@ ) +_environment_adapter: TypeAdapter[Environment] = TypeAdapter(Environment) +_environment_api_key_adapter: TypeAdapter[EnvironmentAPIKey] = TypeAdapter( + EnvironmentAPIKey +) _environment_compressed_adapter: TypeAdapter[EnvironmentCompressed] = TypeAdapter( EnvironmentCompressed, ) @@ -57,19 +64,16 @@ def map_environment_to_environment_document( - environment: "Environment", + environment: "EnvironmentModel", ) -> Document: - return { - field_name: _map_value_to_document_value(value) - for field_name, value in map_environment_to_engine( - environment, - with_integrations=True, - ) - } + return cast( + Document, + _environment_adapter.validate_python(map_environment_to_engine(environment)), + ) def map_environment_to_compressed_environment_document( - environment: "Environment", + environment: "EnvironmentModel", ) -> CompressedEnvironmentDocument: return _get_compressed_environment_document( document=map_environment_to_environment_document(environment), @@ -78,7 +82,7 @@ def map_environment_to_compressed_environment_document( def map_environment_to_environment_v2_document( - environment: "Environment", + environment: "EnvironmentModel", ) -> Document: environment_document = map_environment_to_environment_document(environment) environment_api_key = environment_document.pop("api_key") @@ -91,7 +95,7 @@ def map_environment_to_environment_v2_document( def map_environment_to_compressed_environment_v2_document( - environment: "Environment", + environment: "EnvironmentModel", ) -> CompressedEnvironmentDocument: return _get_compressed_environment_document( document=map_environment_to_environment_v2_document(environment), @@ -100,12 +104,14 @@ def map_environment_to_compressed_environment_v2_document( def map_environment_api_key_to_environment_api_key_document( - environment_api_key: "EnvironmentAPIKey", + environment_api_key: "EnvironmentAPIKeyModel", ) -> Document: - return { - field_name: _map_value_to_document_value(value) - for field_name, value in map_environment_api_key_to_engine(environment_api_key) - } + return cast( + Document, + _environment_api_key_adapter.validate_python( + map_environment_api_key_to_engine(environment_api_key) + ), + ) def map_engine_identity_to_identity_document( diff --git a/api/util/mappers/engine.py b/api/util/mappers/engine.py index 77ffc7a6bb02..b3f07d28b7ec 100644 --- a/api/util/mappers/engine.py +++ b/api/util/mappers/engine.py @@ -1,32 +1,11 @@ from collections.abc import Iterable from itertools import chain -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional from uuid import UUID from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES from features.versioning.models import EnvironmentFeatureVersion -from util.engine_models.environments.integrations.models import IntegrationModel -from util.engine_models.environments.models import ( - EnvironmentAPIKeyModel, - EnvironmentModel, - WebhookModel, -) -from util.engine_models.features.models import ( - FeatureModel, - FeatureSegmentModel, - FeatureStateModel, - MultivariateFeatureOptionModel, - MultivariateFeatureStateValueModel, -) from util.engine_models.identities.models import IdentityModel -from util.engine_models.identities.traits.models import TraitModel -from util.engine_models.organisations.models import OrganisationModel -from util.engine_models.projects.models import ProjectModel -from util.engine_models.segments.models import ( - SegmentConditionModel, - SegmentModel, - SegmentRuleModel, -) if TYPE_CHECKING: # pragma: no cover from environments.identities.models import ( # type: ignore[attr-defined] @@ -57,73 +36,71 @@ ) -def map_traits_to_engine(traits: Iterable["Trait"]) -> list[TraitModel]: +def map_traits_to_engine(traits: Iterable["Trait"]) -> list[dict[str, Any]]: return [ - TraitModel(trait_key=trait.trait_key, trait_value=trait.trait_value) + {"trait_key": trait.trait_key, "trait_value": trait.trait_value} for trait in traits ] def map_segment_to_engine( segment: "Segment", -) -> SegmentModel: +) -> dict[str, Any]: segment_rules = segment.rules.all() # No reading from ORM past this point! - return SegmentModel( - id=segment.pk, - name=segment.name, - rules=[ + return { + "id": segment.pk, + "name": segment.name, + "rules": [ map_segment_rule_to_engine(segment_rule) for segment_rule in segment_rules ], - ) + "feature_states": [], + } def map_segment_rule_to_engine( segment_rule: "SegmentRule", -) -> SegmentRuleModel: +) -> dict[str, Any]: segment_sub_rules = segment_rule.rules.all() conditions = segment_rule.conditions.all() - return SegmentRuleModel( - type=segment_rule.type, # type: ignore[arg-type] - rules=[ + return { + "type": segment_rule.type, + "rules": [ map_segment_rule_to_engine(segment_sub_rule) for segment_sub_rule in segment_sub_rules ], - conditions=[ - SegmentConditionModel( - operator=condition.operator, # type: ignore[arg-type] - value=condition.value, - property_=condition.property, - ) + "conditions": [ + { + "operator": condition.operator, + "value": condition.value, + "property_": condition.property, + } for condition in conditions ], - ) + } def map_integration_to_engine( integration: Optional["EnvironmentIntegrationModel"], -) -> Optional[IntegrationModel]: +) -> Optional[dict[str, Any]]: if not integration: return None - return IntegrationModel( - api_key=integration.api_key, - base_url=integration.base_url, - entity_selector=getattr(integration, "entity_selector", None), - ) + return { + "api_key": integration.api_key, + "base_url": integration.base_url, + "entity_selector": getattr(integration, "entity_selector", None), + } def map_webhook_config_to_engine( webhook_config: Optional["WebhookConfiguration"], -) -> Optional[WebhookModel]: +) -> Optional[dict[str, Any]]: if not webhook_config: return None - return WebhookModel( - url=webhook_config.url, - secret=webhook_config.secret, - ) + return {"url": webhook_config.url, "secret": webhook_config.secret} def map_feature_state_to_engine( @@ -131,67 +108,60 @@ def map_feature_state_to_engine( *, mv_fs_values: Optional[Iterable["MultivariateFeatureStateValue"]] = None, metadata: Optional[dict[str, object]] = None, -) -> FeatureStateModel: +) -> dict[str, Any]: feature = feature_state.feature feature_segment: Optional["FeatureSegment"] = feature_state.feature_segment - if feature_segment: - feature_segment_model = FeatureSegmentModel( - priority=feature_segment.priority, - ) - else: - feature_segment_model = None - - return FeatureStateModel( - metadata=metadata, - enabled=feature_state.enabled, + return { + "feature": map_feature_to_engine(feature), + "enabled": feature_state.enabled, # The engine and SDKs seed multivariate variant allocation on django_id, # so feeding it the bucketing seed keeps variant assignment stable when - # a feature state is recreated, without changing the engine model or - # environment document schema. See issue #7913. - django_id=feature_state.mv_hashing_seed, - feature_state_value=feature_state.get_feature_state_value(), - featurestate_uuid=feature_state.uuid, - feature_segment=feature_segment_model, - feature=map_feature_to_engine(feature), - multivariate_feature_state_values=[ # type: ignore[arg-type] + # a feature state is recreated, without changing the environment + # document schema. See issue #7913. + "django_id": feature_state.mv_hashing_seed, + "feature_segment": ( + {"priority": feature_segment.priority} if feature_segment else None + ), + "featurestate_uuid": feature_state.uuid, + "feature_state_value": feature_state.get_feature_state_value(), + "multivariate_feature_state_values": [ map_mv_fs_value_to_engine(mv_fs_value) for mv_fs_value in mv_fs_values or [] ], - ) + **({"metadata": metadata} if metadata else {}), + } def map_mv_fs_value_to_engine( mv_fs_value: "MultivariateFeatureStateValue", -) -> MultivariateFeatureStateValueModel: +) -> dict[str, Any]: mv_feature_option: "MultivariateFeatureOption" = ( mv_fs_value.multivariate_feature_option ) - return MultivariateFeatureStateValueModel( - percentage_allocation=mv_fs_value.percentage_allocation, - id=mv_fs_value.id, - mv_fs_value_uuid=mv_fs_value.uuid, - multivariate_feature_option=map_mv_option_to_engine(mv_feature_option), - ) + return { + "multivariate_feature_option": map_mv_option_to_engine(mv_feature_option), + "percentage_allocation": mv_fs_value.percentage_allocation, + "id": mv_fs_value.id, + "mv_fs_value_uuid": mv_fs_value.uuid, + } -def map_feature_to_engine(feature: "Feature") -> FeatureModel: - return FeatureModel(id=feature.pk, name=feature.name, type=feature.type) +def map_feature_to_engine(feature: "Feature") -> dict[str, Any]: + return {"id": feature.pk, "name": feature.name, "type": feature.type} def map_mv_option_to_engine( mv_option: "MultivariateFeatureOption", -) -> MultivariateFeatureOptionModel: - return MultivariateFeatureOptionModel( - value=mv_option.value, id=mv_option.id, key=mv_option.key - ) +) -> dict[str, Any]: + return {"value": mv_option.value, "id": mv_option.id, "key": mv_option.key} def map_environment_to_engine( environment: "Environment", *, with_integrations: bool = True, -) -> EnvironmentModel: +) -> dict[str, Any]: """ Maps Core API's `environments.models.Environment` model instance to the flag_engine environment document. @@ -199,7 +169,6 @@ def map_environment_to_engine( feature versions. :param Environment environment: the environment to map - :rtype EnvironmentModel """ from experimentation.feature_state_metadata import ( # avoid circular import get_feature_state_metadata_builder, @@ -273,22 +242,22 @@ def map_environment_to_engine( # No reading from ORM past this point! # Prepare relationships. - organisation_model = OrganisationModel( - id=organisation.pk, - name=organisation.name, - feature_analytics=organisation.feature_analytics, - stop_serving_flags=organisation.stop_serving_flags, - persist_trait_data=organisation.persist_trait_data, - ) + organisation_model = { + "id": organisation.pk, + "name": organisation.name, + "feature_analytics": organisation.feature_analytics, + "stop_serving_flags": organisation.stop_serving_flags, + "persist_trait_data": organisation.persist_trait_data, + } project_segment_models = [ - SegmentModel( - id=segment.pk, - name=segment.name, - rules=[ + { + "id": segment.pk, + "name": segment.name, + "rules": [ map_segment_rule_to_engine(segment_rule) for segment_rule in project_segment_rules_by_segment_id.pop(segment.pk) ], - feature_states=[ + "feature_states": [ map_feature_state_to_engine( feature_state, mv_fs_values=multivariate_feature_state_values_by_feature_state_id.pop( @@ -300,22 +269,22 @@ def map_environment_to_engine( segment.pk ) ], - ) + } for segment in project_segments ] - project_model = ProjectModel( - id=project.pk, - name=project.name, - hide_disabled_flags=project.hide_disabled_flags, - enable_realtime_updates=project.enable_realtime_updates, - server_key_only_feature_ids=[ + project_model = { + "id": project.pk, + "name": project.name, + "organisation": organisation_model, + "hide_disabled_flags": project.hide_disabled_flags, + "segments": project_segment_models, + "enable_realtime_updates": project.enable_realtime_updates, + "server_key_only_feature_ids": [ feature.pk for feature_state in environment_feature_states if (feature := feature_state.feature).is_server_key_only ], - organisation=organisation_model, - segments=project_segment_models, - ) + } feature_state_models = [ map_feature_state_to_engine( feature_state, @@ -347,48 +316,48 @@ def map_environment_to_engine( integration_configs.pop("webhook_config", None), ) - return EnvironmentModel( + return { # # Attributes: - id=environment.pk, - api_key=environment.api_key, - name=environment.name, - allow_client_traits=environment.allow_client_traits, - updated_at=environment.updated_at, - use_identity_composite_key_for_hashing=environment.use_identity_composite_key_for_hashing, - hide_sensitive_data=environment.hide_sensitive_data, - hide_disabled_flags=environment.hide_disabled_flags, - use_identity_overrides_in_local_eval=environment.use_identity_overrides_in_local_eval, - onboarding_pending=environment.first_evaluated_at is None, + "id": environment.pk, + "api_key": environment.api_key, + "name": environment.name, + "allow_client_traits": environment.allow_client_traits, + "updated_at": environment.updated_at, + "hide_sensitive_data": environment.hide_sensitive_data, + "hide_disabled_flags": environment.hide_disabled_flags, + "use_identity_composite_key_for_hashing": environment.use_identity_composite_key_for_hashing, + "use_identity_overrides_in_local_eval": environment.use_identity_overrides_in_local_eval, + "onboarding_pending": environment.first_evaluated_at is None, # # Relationships: - project=project_model, - feature_states=feature_state_models, + "project": project_model, + "feature_states": feature_state_models, + "identity_overrides": [], # # Integrations: - amplitude_config=amplitude_config_model, - heap_config=heap_config_model, - mixpanel_config=mixpanel_config_model, - rudderstack_config=rudderstack_config_model, - segment_config=segment_config_model, - webhook_config=webhook_config_model, - ) + "amplitude_config": amplitude_config_model, + "dynatrace_config": None, + "heap_config": heap_config_model, + "mixpanel_config": mixpanel_config_model, + "rudderstack_config": rudderstack_config_model, + "segment_config": segment_config_model, + "webhook_config": webhook_config_model, + } def map_environment_api_key_to_engine( environment_api_key: "EnvironmentAPIKey", -) -> EnvironmentAPIKeyModel: - client_api_key = environment_api_key.environment.api_key - - return EnvironmentAPIKeyModel( - id=environment_api_key.pk, - key=environment_api_key.key, - created_at=environment_api_key.created_at, - name=environment_api_key.name, - client_api_key=client_api_key, - expires_at=environment_api_key.expires_at, - active=environment_api_key.active, - ) +) -> dict[str, Any]: + return { + "id": environment_api_key.pk, + "key": environment_api_key.key, + "created_at": environment_api_key.created_at, + "name": environment_api_key.name, + "client_api_key": environment_api_key.environment.api_key, + "expires_at": environment_api_key.expires_at, + "active": environment_api_key.active, + } def map_identity_to_engine( @@ -428,16 +397,18 @@ def map_identity_to_engine( ] identity_trait_models = map_traits_to_engine(identity_traits) - return IdentityModel( - # Attributes: - identifier=identity.identifier, - environment_api_key=environment_api_key, - created_date=identity.created_date, - django_id=identity.pk, - # - # Relationships: - identity_features=identity_feature_state_models, # type: ignore[arg-type] - identity_traits=identity_trait_models, + return IdentityModel.model_validate( + { + # Attributes: + "identifier": identity.identifier, + "environment_api_key": environment_api_key, + "created_date": identity.created_date, + "django_id": identity.pk, + # + # Relationships: + "identity_features": identity_feature_state_models, + "identity_traits": identity_trait_models, + } ) diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py index 26c23ec214f9..80b82e2aea29 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -15,12 +15,10 @@ ) SDKDocument: TypeAlias = dict[str, SDKDocumentValue] -SDK_DOCUMENT_EXCLUDE: dict[str, bool | dict[str, set[str]]] = { - **dict.fromkeys(IDENTITY_INTEGRATIONS_RELATION_NAMES, True), - "dynatrace_config": True, - "onboarding_pending": True, - # System-owned identity data must never reach local-eval SDKs. - "identity_overrides": {"__all__": {"system_traits"}}, +SDK_DOCUMENT_EXCLUDE = { + *IDENTITY_INTEGRATIONS_RELATION_NAMES, + "dynatrace_config", + "onboarding_pending", } @@ -39,9 +37,16 @@ def map_environment_to_sdk_document(environment: "Environment") -> SDKDocument: identity_id not in identities_with_overrides ): identities_with_overrides[identity_id] = feature_state.identity - engine_environment.identity_overrides = [ - map_identity_to_engine(identity, with_traits=False) + engine_environment["identity_overrides"] = [ + # System-owned identity data must never reach local-eval SDKs. + map_identity_to_engine(identity, with_traits=False).model_dump( + exclude={"system_traits"} + ) for identity in identities_with_overrides.values() ] - return engine_environment.model_dump(exclude=SDK_DOCUMENT_EXCLUDE) + return { + key: value + for key, value in engine_environment.items() + if key not in SDK_DOCUMENT_EXCLUDE + } From aac37fe21e33899bec19b31b51097e90e281df2c Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Fri, 25 Sep 2026 22:36:26 +0200 Subject: [PATCH 3/5] refactor(api): drop the vendored engine models for flagsmith_schemas documents Edge identities and identity overrides are now plain flagsmith_schemas TypedDicts, validated through TypeAdapters. Stored numbers are validated as native numbers, since the schema doesn't round-trip Decimals. --- api/app/pagination.py | 4 +- api/e2etests/e2e_seed_data.py | 10 +- .../identities/edge_identity_service.py | 5 +- api/edge_api/identities/exceptions.py | 4 + api/edge_api/identities/export.py | 4 +- api/edge_api/identities/models.py | 194 ++++++++++++------ api/edge_api/identities/serializers.py | 106 +++++----- api/edge_api/identities/utils.py | 11 +- api/edge_api/identities/views.py | 50 +++-- api/environments/dynamodb/services.py | 8 +- api/environments/dynamodb/types.py | 14 +- .../dynamodb/wrappers/environment_wrapper.py | 8 +- .../dynamodb/wrappers/identity_wrapper.py | 16 +- api/environments/identities/models.py | 4 +- api/environments/identities/serializers.py | 10 +- api/evaluation/mappers.py | 60 +++--- api/evaluation/services.py | 7 +- api/evaluation/types.py | 5 +- api/features/types.py | 5 +- api/metrics/metrics_service.py | 2 +- api/tests/conftest.py | 18 +- .../edge_api/identities/conftest.py | 5 +- ...est_edge_identity_featurestates_viewset.py | 66 ++---- .../test_edge_api_identities_serializers.py | 9 +- .../identities/test_edge_identity_models.py | 89 ++++---- .../test_unit_edge_api_identities_tasks.py | 11 +- .../dynamodb/test_unit_services.py | 6 +- ...st_unit_dynamodb_environment_v2_wrapper.py | 58 +++--- .../test_unit_evaluation_mappers.py | 67 +++--- .../test_unit_evaluation_services.py | 89 +++----- .../test_unit_features_features_service.py | 27 ++- .../unit/features/test_unit_features_views.py | 6 +- .../unit/metrics/test_unit_metrics_service.py | 4 +- .../identities/test_unit_identities_models.py | 31 --- .../traits/test_unit_traits_types.py | 31 --- .../mappers/test_unit_mappers_dynamodb.py | 93 ++++++--- .../util/mappers/test_unit_mappers_engine.py | 80 +++----- .../util/mappers/test_unit_mappers_sdk.py | 4 +- api/util/engine_models/__init__.py | 8 - .../engine_models/environments/__init__.py | 0 .../environments/integrations/__init__.py | 0 .../environments/integrations/models.py | 9 - api/util/engine_models/environments/models.py | 51 ----- api/util/engine_models/features/__init__.py | 0 api/util/engine_models/features/models.py | 84 -------- api/util/engine_models/identities/__init__.py | 0 api/util/engine_models/identities/models.py | 96 --------- .../identities/traits/__init__.py | 0 .../identities/traits/constants.py | 1 - .../engine_models/identities/traits/models.py | 8 - .../engine_models/identities/traits/types.py | 62 ------ .../engine_models/organisations/__init__.py | 0 .../engine_models/organisations/models.py | 9 - api/util/engine_models/projects/__init__.py | 0 api/util/engine_models/projects/models.py | 16 -- api/util/engine_models/segments/__init__.py | 0 api/util/engine_models/segments/models.py | 28 --- api/util/engine_models/utils/__init__.py | 0 api/util/engine_models/utils/datetime.py | 5 - api/util/engine_models/utils/exceptions.py | 6 - api/util/mappers/__init__.py | 6 + api/util/mappers/dynamodb.py | 157 +++++++------- api/util/mappers/engine.py | 47 +++-- api/util/mappers/sdk.py | 10 +- 64 files changed, 710 insertions(+), 1114 deletions(-) delete mode 100644 api/tests/unit/util/engine_models/identities/test_unit_identities_models.py delete mode 100644 api/tests/unit/util/engine_models/identities/traits/test_unit_traits_types.py delete mode 100644 api/util/engine_models/__init__.py delete mode 100644 api/util/engine_models/environments/__init__.py delete mode 100644 api/util/engine_models/environments/integrations/__init__.py delete mode 100644 api/util/engine_models/environments/integrations/models.py delete mode 100644 api/util/engine_models/environments/models.py delete mode 100644 api/util/engine_models/features/__init__.py delete mode 100644 api/util/engine_models/features/models.py delete mode 100644 api/util/engine_models/identities/__init__.py delete mode 100644 api/util/engine_models/identities/models.py delete mode 100644 api/util/engine_models/identities/traits/__init__.py delete mode 100644 api/util/engine_models/identities/traits/constants.py delete mode 100644 api/util/engine_models/identities/traits/models.py delete mode 100644 api/util/engine_models/identities/traits/types.py delete mode 100644 api/util/engine_models/organisations/__init__.py delete mode 100644 api/util/engine_models/organisations/models.py delete mode 100644 api/util/engine_models/projects/__init__.py delete mode 100644 api/util/engine_models/projects/models.py delete mode 100644 api/util/engine_models/segments/__init__.py delete mode 100644 api/util/engine_models/segments/models.py delete mode 100644 api/util/engine_models/utils/__init__.py delete mode 100644 api/util/engine_models/utils/datetime.py delete mode 100644 api/util/engine_models/utils/exceptions.py diff --git a/api/app/pagination.py b/api/app/pagination.py index 3a31437c2906..8837a5a47ebf 100644 --- a/api/app/pagination.py +++ b/api/app/pagination.py @@ -6,7 +6,7 @@ from rest_framework.pagination import PageNumberPagination from rest_framework.response import Response -from util.engine_models.identities.models import IdentityModel +from edge_api.identities.models import EdgeIdentity class CustomPagination(PageNumberPagination): @@ -27,7 +27,7 @@ def paginate_queryset(self, dynamo_queryset, request, view=None): # type: ignor ) return [ - IdentityModel.model_validate(identity_document) + EdgeIdentity.from_identity_document(identity_document) for identity_document in dynamo_queryset["Items"] ] diff --git a/api/e2etests/e2e_seed_data.py b/api/e2etests/e2e_seed_data.py index f00717c897de..549a9dad210e 100644 --- a/api/e2etests/e2e_seed_data.py +++ b/api/e2etests/e2e_seed_data.py @@ -24,7 +24,6 @@ from organisations.subscriptions.constants import ENTERPRISE from projects.models import Project, UserProjectPermission from users.models import FFAdminUser, UserPermissionGroup -from util.engine_models.identities.models import IdentityModel as EngineIdentity # Password used by all the test users PASSWORD = "Str0ngp4ssw0rd!" @@ -208,10 +207,9 @@ def seed_data() -> None: for identity_info in identities_test_data: if settings.IDENTITIES_TABLE_NAME_DYNAMO: - engine_identity = EngineIdentity( # pragma: no cover - identifier=identity_info["identifier"], - environment_api_key=identity_info["environment"].api_key, - ) - EdgeIdentity(engine_identity).save() # pragma: no cover + EdgeIdentity.create( # pragma: no cover + identity_info["identifier"], + identity_info["environment"].api_key, + ).save() else: Identity.objects.create(**identity_info) diff --git a/api/edge_api/identities/edge_identity_service.py b/api/edge_api/identities/edge_identity_service.py index 5fda5c4f4d87..9849fddc4584 100644 --- a/api/edge_api/identities/edge_identity_service.py +++ b/api/edge_api/identities/edge_identity_service.py @@ -5,6 +5,7 @@ from environments.dynamodb.types import ( IdentityOverrideV2, ) +from util.mappers import map_identity_override_document_to_identity_override ddb_environment_v2_wrapper = DynamoEnvironmentV2Wrapper() @@ -20,7 +21,7 @@ def get_edge_identity_overrides( ) ) return [ - IdentityOverrideV2.model_validate( + map_identity_override_document_to_identity_override( {**item, "environment_id": str(item["environment_id"])} ) for item in override_items @@ -49,4 +50,4 @@ def get_overridden_feature_ids_for_edge_identity(identity_uuid: str) -> set[int] except ObjectDoesNotExist: return set() identity = EdgeIdentity.from_identity_document(identity_document) - return {fs.feature.id for fs in identity.feature_overrides} + return {int(fs["feature"]["id"]) for fs in identity.feature_overrides} diff --git a/api/edge_api/identities/exceptions.py b/api/edge_api/identities/exceptions.py index 9f25db3c135d..50a9ac35dff5 100644 --- a/api/edge_api/identities/exceptions.py +++ b/api/edge_api/identities/exceptions.py @@ -4,3 +4,7 @@ class TraitPersistenceError(APIException): status_code = status.HTTP_400_BAD_REQUEST + + +class DuplicateFeatureState(ValueError): + pass diff --git a/api/edge_api/identities/export.py b/api/edge_api/identities/export.py index b0de70c9ab9b..70302502f8b5 100644 --- a/api/edge_api/identities/export.py +++ b/api/edge_api/identities/export.py @@ -4,12 +4,12 @@ from decimal import Decimal from django.utils import timezone +from flag_engine.context.mappers import map_any_value_to_context_value from edge_api.identities.models import EdgeIdentity from environments.identities.traits.models import Trait from features.models import Feature, FeatureState from features.multivariate.models import MultivariateFeatureOption -from util.engine_models.identities.traits.types import map_any_value_to_trait_value EXPORT_EDGE_IDENTITY_PAGINATION_LIMIT = 20000 @@ -122,7 +122,7 @@ def get_mv_feature_option_uuid_cache(environment_api_key: str) -> dict[int, str] def export_edge_trait(trait: dict, identifier: str, environment_api_key: str) -> dict: # type: ignore[type-arg] - trait_value = map_any_value_to_trait_value(trait["trait_value"]) + trait_value = map_any_value_to_context_value(trait["trait_value"]) trait_value_data = Trait.generate_trait_value_data(trait_value) return { "model": "traits.trait", diff --git a/api/edge_api/identities/models.py b/api/edge_api/identities/models.py index e8430e3920a3..bff24ed9f710 100644 --- a/api/edge_api/identities/models.py +++ b/api/edge_api/identities/models.py @@ -1,12 +1,16 @@ import copy import typing +import uuid from contextlib import suppress from datetime import timedelta from django.conf import settings from django.utils import timezone +from flag_engine.context.mappers import map_any_value_to_context_value +from flagsmith_schemas.dynamodb import FeatureState, Identity from api_keys.user import APIKeyUser +from edge_api.identities.exceptions import DuplicateFeatureState from edge_api.identities.tasks import ( generate_audit_log_records, sync_identity_document_features, @@ -18,43 +22,81 @@ from environments.models import Environment from evaluation.services import get_edge_identity_override_value from users.models import FFAdminUser -from util.engine_models.features.models import FeatureStateModel -from util.engine_models.identities.models import IdentityFeaturesList, IdentityModel -from util.mappers import map_engine_identity_to_identity_document +from util.mappers import ( + map_engine_identity_to_identity_document, + map_identifier_to_engine, + map_identity_document_to_engine_identity, +) + + +def new_feature_override(**fields: typing.Any) -> FeatureState: + """An identity override's fields, defaulted as for a new one.""" + return typing.cast( + FeatureState, + { + "django_id": None, + "feature_segment": None, + "featurestate_uuid": str(uuid.uuid4()), + "feature_state_value": None, + **fields, + "multivariate_feature_state_values": [ + {"mv_fs_value_uuid": str(uuid.uuid4()), **mv_fs_value} + for mv_fs_value in fields.get("multivariate_feature_state_values", []) + ], + }, + ) class EdgeIdentity: dynamo_wrapper = DynamoIdentityWrapper() - def __init__(self, engine_identity_model: IdentityModel): - self.engine_identity_model = engine_identity_model + def __init__(self, document: Identity): + document.setdefault("identity_features", []) + self.document = document self._reset_initial_state() # type: ignore[no-untyped-call] @classmethod - def from_identity_document(cls, identity_document: dict) -> "EdgeIdentity": # type: ignore[type-arg] - return EdgeIdentity(IdentityModel.model_validate(identity_document)) + def create( + cls, identifier: str, environment_api_key: str, **fields: typing.Any + ) -> "EdgeIdentity": + return cls.from_identity_document( + { + "identifier": identifier, + "environment_api_key": environment_api_key, + **fields, + } + ) + + @classmethod + def from_identity_document( + cls, identity_document: typing.Mapping[str, typing.Any] + ) -> "EdgeIdentity": + return EdgeIdentity( + map_identity_document_to_engine_identity( + map_identifier_to_engine(**identity_document) + ) + ) @property def environment_api_key(self) -> str: - return self.engine_identity_model.environment_api_key + return self.document["environment_api_key"] @property - def feature_overrides(self) -> IdentityFeaturesList: - return self.engine_identity_model.identity_features + def feature_overrides(self) -> list[FeatureState]: + return self.document["identity_features"] @property def id(self) -> typing.Union[int, str]: - return self.engine_identity_model.django_id or str( - self.engine_identity_model.identity_uuid - ) + django_id = self.document.get("django_id") + return int(django_id) if django_id else self.identity_uuid @property def identifier(self) -> str: - return self.engine_identity_model.identifier + return self.document["identifier"] @property def identity_uuid(self) -> str: - return str(self.engine_identity_model.identity_uuid) + return self.document["identity_uuid"] @property def environment(self) -> Environment: @@ -65,52 +107,77 @@ def environment(self) -> Environment: @property def dashboard_alias(self) -> str | None: - return self.engine_identity_model.dashboard_alias + return self.document.get("dashboard_alias") @dashboard_alias.setter def dashboard_alias(self, dashboard_alias: str) -> None: - self.engine_identity_model.dashboard_alias = dashboard_alias + self.document["dashboard_alias"] = dashboard_alias - def add_feature_override(self, feature_state: FeatureStateModel) -> None: - self.engine_identity_model.identity_features.append(feature_state) + def add_feature_override(self, feature_state: FeatureState) -> None: + feature_id = feature_state["feature"]["id"] + if any(fs["feature"]["id"] == feature_id for fs in self.feature_overrides): + raise DuplicateFeatureState( + f"Feature state for feature id={feature_id} already exists" + ) + self.feature_overrides.append(feature_state) def get_feature_state_by_feature_name_or_id( self, feature: typing.Union[str, int] - ) -> typing.Optional[FeatureStateModel]: - def match_feature_state(fs): # type: ignore[no-untyped-def] - if isinstance(feature, int): - return fs.feature.id == feature - return fs.feature.name == feature - - feature_state = next( - filter( - match_feature_state, - self.engine_identity_model.identity_features, - ), + ) -> typing.Optional[FeatureState]: + key = "id" if isinstance(feature, int) else "name" + return next( + (fs for fs in self.feature_overrides if fs["feature"][key] == feature), # type: ignore[literal-required] None, ) - return feature_state - def get_feature_state_by_featurestate_uuid( self, featurestate_uuid: str - ) -> typing.Optional[FeatureStateModel]: + ) -> typing.Optional[FeatureState]: return next( - filter( - lambda fs: str(fs.featurestate_uuid) == featurestate_uuid, # type: ignore[arg-type,union-attr] - self.engine_identity_model.identity_features, + ( + fs + for fs in self.feature_overrides + if str(fs.get("featurestate_uuid")) == featurestate_uuid ), None, ) def get_hash_key(self, use_identity_composite_key_for_hashing: bool) -> str: - return self.engine_identity_model.get_hash_key( - use_identity_composite_key_for_hashing - ) + if use_identity_composite_key_for_hashing: + return self.document["composite_key"] + if (django_id := self.document.get("django_id")) is not None: + return str(django_id) + return self.identifier + + def update_traits( + self, traits: typing.Iterable[typing.Mapping[str, typing.Any]] + ) -> bool: + """Set, or unset if their value is `None`, traits; return whether any changed.""" + existing_traits = { + trait["trait_key"]: trait for trait in self.document["identity_traits"] + } + traits_changed = False + for trait in traits: + trait_key, trait_value = trait["trait_key"], trait["trait_value"] + existing_trait = existing_traits.get(trait_key) + if trait_value is None: + traits_changed |= existing_traits.pop(trait_key, None) is not None + elif ( + existing_trait is None + or map_any_value_to_context_value(existing_trait["trait_value"]) + != trait_value + ): + existing_traits[trait_key] = { + "trait_key": trait_key, + "trait_value": trait_value, + } + traits_changed = True + self.document["identity_traits"] = list(existing_traits.values()) + return traits_changed - def remove_feature_override(self, feature_state: FeatureStateModel) -> None: + def remove_feature_override(self, feature_state: FeatureState) -> None: with suppress(ValueError): # ignore if feature state didn't exist - self.engine_identity_model.identity_features.remove(feature_state) + self.feature_overrides.remove(feature_state) def save(self, user: FFAdminUser | APIKeyUser = None): # type: ignore[no-untyped-def,assignment] self.dynamo_wrapper.put_item(self.to_document()) @@ -122,8 +189,8 @@ def save(self, user: FFAdminUser | APIKeyUser = None): # type: ignore[no-untype self._reset_initial_state() # type: ignore[no-untyped-call] def delete(self, user: FFAdminUser | APIKeyUser = None) -> None: # type: ignore[assignment] - self.dynamo_wrapper.delete_item(self.engine_identity_model.composite_key) - self.engine_identity_model.identity_features.clear() + self.dynamo_wrapper.delete_item(self.document["composite_key"]) + self.feature_overrides.clear() changeset = self._get_changes() self._update_feature_overrides( changeset=changeset, @@ -146,20 +213,23 @@ def delete(self, user: FFAdminUser | APIKeyUser = None) -> None: # type: ignore def synchronise_features(self, valid_feature_names: typing.Collection[str]) -> None: identity_feature_names = { - fs.feature.name for fs in self.engine_identity_model.identity_features + fs["feature"]["name"] for fs in self.feature_overrides } if not identity_feature_names.issubset(valid_feature_names): - self.engine_identity_model.prune_features(list(valid_feature_names)) + self.document["identity_features"] = [ + fs + for fs in self.feature_overrides + if fs["feature"]["name"] in valid_feature_names + ] sync_identity_document_features.delay(args=(str(self.identity_uuid),)) - def to_document(self) -> dict: # type: ignore[type-arg] - return map_engine_identity_to_identity_document(self.engine_identity_model) + def to_document(self) -> dict[str, typing.Any]: + return map_engine_identity_to_identity_document(self.document) def _update_feature_overrides( self, changeset: IdentityChangeset, user: FFAdminUser | APIKeyUser ) -> None: if changeset["feature_overrides"]: - # TODO: would this be simpler if we put a wrapper around FeatureStateModel instead? kwargs = { "environment_api_key": self.environment_api_key, "identifier": self.identifier, @@ -189,10 +259,11 @@ def _get_changes(self) -> IdentityChangeset: changes = {} # type: ignore[var-annotated] feature_changes = changes.setdefault("feature_overrides", {}) previous_feature_overrides = { - fs.featurestate_uuid: fs for fs in previous_instance.feature_overrides + fs.get("featurestate_uuid"): fs + for fs in previous_instance.feature_overrides } current_feature_overrides = { - fs.featurestate_uuid: fs for fs in self.feature_overrides + fs.get("featurestate_uuid"): fs for fs in self.feature_overrides } environment = Environment.get_from_cache(self.environment_api_key) assert environment @@ -200,13 +271,13 @@ def _get_changes(self) -> IdentityChangeset: for uuid_, previous_fs in previous_feature_overrides.items(): current_matching_fs = current_feature_overrides.get(uuid_) if current_matching_fs is None: - feature_changes[previous_fs.feature.name] = generate_change_dict( + feature_changes[previous_fs["feature"]["name"]] = generate_change_dict( change_type="-", edge_identity=self, environment=environment, old=previous_fs, ) - elif current_matching_fs.enabled != previous_fs.enabled or ( + elif current_matching_fs["enabled"] != previous_fs["enabled"] or ( get_edge_identity_override_value( self, current_matching_fs, environment=environment ) @@ -214,7 +285,7 @@ def _get_changes(self) -> IdentityChangeset: self, previous_fs, environment=environment ) ): - feature_changes[previous_fs.feature.name] = generate_change_dict( + feature_changes[previous_fs["feature"]["name"]] = generate_change_dict( change_type="~", edge_identity=self, environment=environment, @@ -224,7 +295,7 @@ def _get_changes(self) -> IdentityChangeset: for uuid_, previous_fs in current_feature_overrides.items(): if uuid_ not in previous_feature_overrides: - feature_changes[previous_fs.feature.name] = generate_change_dict( + feature_changes[previous_fs["feature"]["name"]] = generate_change_dict( change_type="+", edge_identity=self, environment=environment, @@ -246,10 +317,13 @@ def clone_flag_states_from(self, source_identity: "EdgeIdentity") -> None: # Clone identity_source's feature states to identity_target for feature_in_source in source_identity.feature_overrides: - feature_state_target = FeatureStateModel( - feature=feature_in_source.feature, - feature_state_value=feature_in_source.feature_state_value, - enabled=feature_in_source.enabled, - multivariate_feature_state_values=feature_in_source.multivariate_feature_state_values, + self.add_feature_override( + new_feature_override( + feature=feature_in_source["feature"], + feature_state_value=feature_in_source["feature_state_value"], + enabled=feature_in_source["enabled"], + multivariate_feature_state_values=copy.deepcopy( + feature_in_source.get("multivariate_feature_state_values", []) + ), + ) ) - self.add_feature_override(feature_state_target) diff --git a/api/edge_api/identities/serializers.py b/api/edge_api/identities/serializers.py index 75446f6a2804..834b9a975bae 100644 --- a/api/edge_api/identities/serializers.py +++ b/api/edge_api/identities/serializers.py @@ -1,13 +1,28 @@ import copy import typing +import uuid from django.utils import timezone from drf_spectacular.utils import extend_schema_field +from flagsmith_schemas.dynamodb import ( + Feature as EdgeFeature, +) +from flagsmith_schemas.dynamodb import ( + FeatureState as EdgeFeatureState, +) +from flagsmith_schemas.dynamodb import ( + MultivariateFeatureOption as EdgeMultivariateFeatureOption, +) +from flagsmith_schemas.dynamodb import ( + MultivariateFeatureStateValue as EdgeMultivariateFeatureStateValue, +) +from pydantic import TypeAdapter from pydantic import ValidationError as PydanticValidationError from pyngo import drf_error_details from rest_framework import serializers from rest_framework.exceptions import ValidationError +from edge_api.identities.exceptions import DuplicateFeatureState from environments.dynamodb.types import IdentityOverrideV2 from environments.models import Environment from evaluation.services import get_edge_identity_override_value @@ -17,26 +32,15 @@ from features.serializers import ( # type: ignore[attr-defined] FeatureStateValueSerializer, ) -from util.engine_models.features.models import FeatureModel as EngineFeatureModel -from util.engine_models.features.models import ( - FeatureStateModel as EngineFeatureStateModel, -) -from util.engine_models.features.models import ( - MultivariateFeatureOptionModel as EngineMultivariateFeatureOptionModel, -) -from util.engine_models.features.models import ( - MultivariateFeatureStateValueModel as EngineMultivariateFeatureStateValueModel, -) -from util.engine_models.identities.models import IdentityModel as EngineIdentity -from util.engine_models.utils.exceptions import DuplicateFeatureState from util.mappers import ( map_engine_identity_to_identity_document, map_feature_to_engine, + map_identifier_to_engine, map_mv_option_to_engine, ) from webhooks.constants import WEBHOOK_DATETIME_FORMAT -from .models import EdgeIdentity +from .models import EdgeIdentity, new_feature_override from .search import ( DASHBOARD_ALIAS_ATTRIBUTE, DASHBOARD_ALIAS_SEARCH_PREFIX, @@ -46,6 +50,8 @@ ) from .tasks import call_environment_webhook_for_feature_state_change +_feature_state_adapter: TypeAdapter[EdgeFeatureState] = TypeAdapter(EdgeFeatureState) + class LowerCaseCharField(serializers.CharField): def to_representation(self, value: typing.Any) -> str: @@ -67,12 +73,12 @@ def create(self, *args, **kwargs): # type: ignore[no-untyped-def] identifier = self.validated_data.get("identifier") dashboard_alias = self.validated_data.get("dashboard_alias") environment_api_key = self.context["view"].kwargs["environment_api_key"] - self.instance = EngineIdentity( - identifier=identifier, - environment_api_key=environment_api_key, + self.instance = map_identifier_to_engine( + identifier, + environment_api_key, dashboard_alias=dashboard_alias, ) - if EdgeIdentity.dynamo_wrapper.get_item(self.instance.composite_key): + if EdgeIdentity.dynamo_wrapper.get_item(self.instance["composite_key"]): raise ValidationError( f"Identity with identifier: {identifier} already exists" ) @@ -102,12 +108,13 @@ class EdgeMultivariateFeatureOptionField(serializers.IntegerField): def to_internal_value( # type: ignore[override] self, data: typing.Any, - ) -> EngineMultivariateFeatureOptionModel: + ) -> EdgeMultivariateFeatureOption: data = super().to_internal_value(data) - return map_mv_option_to_engine(MultivariateFeatureOption.objects.get(id=data)) + return map_mv_option_to_engine(MultivariateFeatureOption.objects.get(id=data)) # type: ignore[return-value] - def to_representation(self, obj): # type: ignore[no-untyped-def] - return obj.id + def to_representation(self, obj: EdgeMultivariateFeatureOption) -> int | None: # type: ignore[override] + option_id = obj.get("id") + return None if option_id is None else int(option_id) class EdgeMultivariateFeatureStateValueSerializer(serializers.Serializer): # type: ignore[type-arg] @@ -116,7 +123,7 @@ class EdgeMultivariateFeatureStateValueSerializer(serializers.Serializer): # ty def to_internal_value(self, data): # type: ignore[no-untyped-def] data = super().to_internal_value(data) - return EngineMultivariateFeatureStateValueModel(**data) + return {"id": None, "mv_fs_value_uuid": str(uuid.uuid4()), **data} @extend_schema_field( @@ -147,18 +154,18 @@ def to_internal_value(self, data): # type: ignore[no-untyped-def] return FeatureStateValue(**feature_state_value_dict).value -class EdgeFeatureField(serializers.Field[EngineFeatureModel, str | int, int, int]): - def to_representation(self, obj: EngineFeatureModel) -> int: - return obj.id +class EdgeFeatureField(serializers.Field[EdgeFeature, str | int, int, int]): + def to_representation(self, obj: EdgeFeature) -> int: + return int(obj["id"]) - def to_internal_value(self, data: str | int) -> EngineFeatureModel: + def to_internal_value(self, data: str | int) -> EdgeFeature: if isinstance(data, int): - return map_feature_to_engine(Feature.objects.get(id=data)) + return map_feature_to_engine(Feature.objects.get(id=data)) # type: ignore[return-value] environment = Environment.objects.get( api_key=self.context["view"].kwargs["environment_api_key"] ) - return map_feature_to_engine( + return map_feature_to_engine( # type: ignore[return-value] Feature.objects.get( name=data, project=environment.project, @@ -178,10 +185,10 @@ class BaseEdgeIdentityFeatureStateSerializer(serializers.Serializer): # type: i featurestate_uuid = serializers.CharField(required=False, read_only=True) def validate_multivariate_feature_state_values( - self, values: list[EngineMultivariateFeatureStateValueModel] - ) -> list[EngineMultivariateFeatureStateValueModel]: + self, values: list[EdgeMultivariateFeatureStateValue] + ) -> list[EdgeMultivariateFeatureStateValue]: validate_identity_override_allocations( - value.percentage_allocation for value in values + float(value["percentage_allocation"]) for value in values ) return values @@ -195,10 +202,7 @@ def save(self, **kwargs): # type: ignore[no-untyped-def] previous_state = copy.deepcopy(self.instance) if not self.instance: - try: - self.instance = EngineFeatureStateModel.parse_obj(self.validated_data) - except PydanticValidationError as exc: - raise ValidationError(drf_error_details(exc)) + self.instance = new_feature_override(**self.validated_data) try: identity.add_feature_override(self.instance) except DuplicateFeatureState as e: @@ -206,14 +210,18 @@ def save(self, **kwargs): # type: ignore[no-untyped-def] "Feature state already exists." ) from e - self.instance.set_value(feature_state_value) - self.instance.enabled = self.validated_data.get( - "enabled", self.instance.enabled + self.instance["feature_state_value"] = feature_state_value + self.instance["enabled"] = self.validated_data.get( + "enabled", self.instance["enabled"] ) - self.instance.multivariate_feature_state_values = self.validated_data.get( + self.instance["multivariate_feature_state_values"] = self.validated_data.get( "multivariate_feature_state_values", - self.instance.multivariate_feature_state_values, + self.instance.get("multivariate_feature_state_values", []), ) + try: + _feature_state_adapter.validate_python(self.instance) + except PydanticValidationError as exc: + raise ValidationError(drf_error_details(exc)) identity.save(user=request.user) @@ -234,14 +242,16 @@ def save(self, **kwargs): # type: ignore[no-untyped-def] # - move this logic to the EdgeIdentity model call_environment_webhook_for_feature_state_change.delay( kwargs={ - "feature_id": self.instance.feature.id, + "feature_id": int(self.instance["feature"]["id"]), "environment_api_key": identity.environment_api_key, "identity_id": identity.id, "identity_identifier": identity.identifier, "changed_by": str(request.user), - "new_enabled_state": self.instance.enabled, + "new_enabled_state": self.instance["enabled"], "new_value": new_value, - "previous_enabled_state": getattr(previous_state, "enabled", None), + "previous_enabled_state": ( + previous_state["enabled"] if previous_state else None + ), "previous_value": previous_value, "timestamp": timezone.now().strftime(WEBHOOK_DATETIME_FORMAT), }, @@ -335,12 +345,10 @@ def to_representation(self, instance: IdentityOverrideV2): # type: ignore[no-un # and make it available to the field class. to_representation seems like the # best place for this since we only care about serialization here (not # deserialization). - self.context["identity"] = EdgeIdentity.from_identity_document( - { - "identifier": instance.identifier, - "identity_uuid": instance.identity_uuid, - "environment_api_key": self.context["environment"].api_key, - } + self.context["identity"] = EdgeIdentity.create( + instance["identifier"], + self.context["environment"].api_key, + identity_uuid=instance["identity_uuid"], ) return super().to_representation(instance) diff --git a/api/edge_api/identities/utils.py b/api/edge_api/identities/utils.py index 442c9d4e82cb..6da9728fc8fc 100644 --- a/api/edge_api/identities/utils.py +++ b/api/edge_api/identities/utils.py @@ -1,9 +1,10 @@ import typing from evaluation.services import get_edge_identity_override_value -from util.engine_models.features.models import FeatureStateModel if typing.TYPE_CHECKING: + from flagsmith_schemas.dynamodb import FeatureState + from edge_api.identities.models import EdgeIdentity from edge_api.identities.types import ChangeType, FeatureStateChangeDetails from environments.models import Environment @@ -14,8 +15,8 @@ def generate_change_dict( *, edge_identity: "EdgeIdentity", environment: "Environment", - new: FeatureStateModel | None = None, - old: FeatureStateModel | None = None, + new: "FeatureState | None" = None, + old: "FeatureState | None" = None, ) -> "FeatureStateChangeDetails": if not (new or old): raise ValueError("Must provide one of 'new' or 'old'") @@ -41,10 +42,10 @@ def _get_overridden_feature_state_dict( *, edge_identity: "EdgeIdentity", environment: "Environment", - feature_state: FeatureStateModel, + feature_state: "FeatureState", ) -> dict[str, typing.Any]: return { - **feature_state.dict(), + **feature_state, "feature_state_value": get_edge_identity_override_value( edge_identity, feature_state, environment=environment ), diff --git a/api/edge_api/identities/views.py b/api/edge_api/identities/views.py index 2c2b855f4987..169df27e13d1 100644 --- a/api/edge_api/identities/views.py +++ b/api/edge_api/identities/views.py @@ -9,6 +9,9 @@ ) from django.shortcuts import get_object_or_404 from drf_spectacular.utils import extend_schema +from flag_engine.context.mappers import map_any_value_to_context_value +from flagsmith_schemas.api import TraitInput +from pydantic import TypeAdapter from pyngo import drf_error_details from rest_framework import status, viewsets from rest_framework.decorators import action, api_view, permission_classes @@ -49,14 +52,13 @@ from environments.identities.serializers import ( IdentityAllFeatureStatesSerializer, ) +from environments.identities.traits.constants import TRAIT_STRING_VALUE_MAX_LENGTH from environments.models import Environment from environments.permissions.permissions import NestedEnvironmentPermissions from evaluation.services import get_edge_identity_feature_states from features.models import FeatureState from features.permissions import IdentityFeatureStatePermissions from projects.exceptions import DynamoNotEnabledError -from util.engine_models.identities.models import IdentityFeaturesList, IdentityModel -from util.engine_models.identities.traits.models import TraitModel from . import edge_identity_service from .exceptions import TraitPersistenceError @@ -67,6 +69,8 @@ ) from .search import EdgeIdentitySearchData +_trait_input_adapter: TypeAdapter[TraitInput] = TypeAdapter(TraitInput) + class EdgeIdentityViewSet( GenericViewSet, # type: ignore[type-arg] @@ -162,8 +166,11 @@ def perform_destroy(self, instance: EdgeIdentity) -> None: def get_traits(self, request, *args, **kwargs): # type: ignore[no-untyped-def] edge_identity = self.get_object() data = [ - trait.dict() - for trait in edge_identity.engine_identity_model.identity_traits + { + "trait_key": trait["trait_key"], + "trait_value": map_any_value_to_context_value(trait["trait_value"]), + } + for trait in edge_identity.document["identity_traits"] ] return Response(data=data, status=status.HTTP_200_OK) @@ -180,17 +187,28 @@ def update_traits(self, request, *args, **kwargs): # type: ignore[no-untyped-de if not isinstance(request.data, dict): raise ValidationError({"detail": "Request data must be a JSON object."}) try: - trait = TraitModel(**request.data) + trait_input = _trait_input_adapter.validate_python(request.data) except pydantic.ValidationError as validation_error: raise ValidationError( drf_error_details(validation_error) ) from validation_error - _, traits_updated = edge_identity.engine_identity_model.update_traits([trait]) - if traits_updated: + trait_value = trait_input["trait_value"] + if trait_value is not None: + trait_value = map_any_value_to_context_value(trait_value) + if ( + isinstance(trait_value, str) + and len(trait_value) > TRAIT_STRING_VALUE_MAX_LENGTH + ): + raise ValidationError( + { + "trait_value": f"Must be at most {TRAIT_STRING_VALUE_MAX_LENGTH} characters." + } + ) + trait = {"trait_key": trait_input["trait_key"], "trait_value": trait_value} + if edge_identity.update_traits([trait]): edge_identity.save() - data = trait.dict() - return Response(data, status=status.HTTP_200_OK) + return Response(trait, status=status.HTTP_200_OK) class EdgeIdentityFeatureStateViewSet(viewsets.ModelViewSet): # type: ignore[type-arg] @@ -263,13 +281,13 @@ def list(self, request, *args, **kwargs): # type: ignore[no-untyped-def] ) q_params_serializer.is_valid(raise_exception=True) - identity_features: IdentityFeaturesList = self.identity.feature_overrides + identity_features = self.identity.feature_overrides feature = q_params_serializer.data.get("feature") if feature: - identity_features = filter( # type: ignore[assignment] - lambda fs: fs.feature.id == feature, identity_features - ) + identity_features = [ + fs for fs in identity_features if fs["feature"]["id"] == feature + ] serializer = self.get_serializer(identity_features, many=True) return Response(data=serializer.data, status=status.HTTP_200_OK) @@ -336,11 +354,7 @@ def initial(self, request, *args, **kwargs): # type: ignore[no-untyped-def] if identity_document: self.identity = EdgeIdentity.from_identity_document(identity_document) else: - self.identity = EdgeIdentity( - engine_identity_model=IdentityModel( - identifier=identifier, environment_api_key=environment_api_key - ) - ) + self.identity = EdgeIdentity.create(identifier, environment_api_key) @extend_schema( request=EdgeIdentityWithIdentifierFeatureStateRequestBody, diff --git a/api/environments/dynamodb/services.py b/api/environments/dynamodb/services.py index 5f613e8290a8..988b3a34f4a1 100644 --- a/api/environments/dynamodb/services.py +++ b/api/environments/dynamodb/services.py @@ -14,7 +14,6 @@ ) from environments.models import Environment from projects.models import EdgeV2MigrationStatus -from util.engine_models.identities.models import IdentityModel from util.mappers import map_engine_feature_state_to_identity_override logger = logging.getLogger(__name__) @@ -94,12 +93,11 @@ def _iter_paginated_overrides( projection_expression="environment_api_key, identifier, identity_features, identity_uuid", overrides_only=True, ): - identity = IdentityModel.model_validate(item) - for feature_state in identity.identity_features: + for feature_state in item.get("identity_features", []): yield map_engine_feature_state_to_identity_override( feature_state=feature_state, - identity_uuid=str(identity.identity_uuid), - identifier=identity.identifier, + identity_uuid=item["identity_uuid"], + identifier=item["identifier"], environment_api_key=environment_api_key, environment_id=str(environment.id), # type: ignore[arg-type] ) diff --git a/api/environments/dynamodb/types.py b/api/environments/dynamodb/types.py index 53790c7295dd..b68c01ccf61f 100644 --- a/api/environments/dynamodb/types.py +++ b/api/environments/dynamodb/types.py @@ -5,10 +5,7 @@ import boto3 from django.conf import settings -from django.utils import timezone -from pydantic import BaseModel, Field - -from util.engine_models.features.models import FeatureStateModel +from flagsmith_schemas.dynamodb import EnvironmentV2IdentityOverride if typing.TYPE_CHECKING: from projects.models import EdgeV2MigrationStatus @@ -83,14 +80,7 @@ def delete(self): # type: ignore[no-untyped-def] project_metadata_table.delete_item(Key={"id": self.id}) -class IdentityOverrideV2(BaseModel): - environment_id: str - document_key: str - environment_api_key: str - identifier: str - identity_uuid: str - feature_state: FeatureStateModel - created_date: datetime = Field(default_factory=timezone.now) +IdentityOverrideV2: typing.TypeAlias = EnvironmentV2IdentityOverride @dataclass diff --git a/api/environments/dynamodb/wrappers/environment_wrapper.py b/api/environments/dynamodb/wrappers/environment_wrapper.py index 0ac4ee81fc08..50e81dfa3fa2 100644 --- a/api/environments/dynamodb/wrappers/environment_wrapper.py +++ b/api/environments/dynamodb/wrappers/environment_wrapper.py @@ -180,8 +180,12 @@ def update_identity_overrides( for identity_override_to_delete in to_delete: writer.delete_item( Key={ - ENVIRONMENTS_V2_PARTITION_KEY: identity_override_to_delete.environment_id, - ENVIRONMENTS_V2_SORT_KEY: identity_override_to_delete.document_key, + ENVIRONMENTS_V2_PARTITION_KEY: identity_override_to_delete[ + "environment_id" + ], + ENVIRONMENTS_V2_SORT_KEY: identity_override_to_delete[ + "document_key" + ], }, ) for identity_override_to_put in to_put: diff --git a/api/environments/dynamodb/wrappers/identity_wrapper.py b/api/environments/dynamodb/wrappers/identity_wrapper.py index 69cd36571555..9b4011070f1e 100644 --- a/api/environments/dynamodb/wrappers/identity_wrapper.py +++ b/api/environments/dynamodb/wrappers/identity_wrapper.py @@ -22,9 +22,9 @@ from environments.identities.traits.constants import ( TRAIT_STRING_VALUE_MAX_LENGTH, ) -from util.engine_models.identities.models import IdentityModel from util.mappers import ( map_engine_identity_to_identity_document, + map_identifier_to_engine, map_identity_to_identity_document, ) @@ -113,9 +113,7 @@ def set_system_trait( "System trait value must be at most " f"{TRAIT_STRING_VALUE_MAX_LENGTH} characters." ) - composite_key = IdentityModel.generate_composite_key( - environment_api_key, identifier - ) + composite_key = f"{environment_api_key}_{identifier}" # DynamoDB rejects floats and returns all numbers as Decimal. document_value: bool | int | Decimal | str = ( Decimal(str(trait_value)) if isinstance(trait_value, float) else trait_value @@ -135,9 +133,9 @@ def set_system_trait( if document is None: self.table.put_item( # type: ignore[union-attr] Item=map_engine_identity_to_identity_document( - IdentityModel( - identifier=identifier, - environment_api_key=environment_api_key, + map_identifier_to_engine( + identifier, + environment_api_key, system_traits={trait_key: trait_value}, ) ), @@ -187,9 +185,7 @@ def unset_system_trait( trait_key: str, ) -> None: """Idempotently remove a system trait from an identity document.""" - composite_key = IdentityModel.generate_composite_key( - environment_api_key, identifier - ) + composite_key = f"{environment_api_key}_{identifier}" try: self.table.update_item( # type: ignore[union-attr] Key={"composite_key": composite_key}, diff --git a/api/environments/identities/models.py b/api/environments/identities/models.py index 61f0211f3b3a..46670e3b11d0 100644 --- a/api/environments/identities/models.py +++ b/api/environments/identities/models.py @@ -107,11 +107,11 @@ def generate_traits( ) -> list[Trait]: """ Given a list of trait data items, validated by TraitSerializerFull, generate - a list of TraitModel objects for the given identity. + a list of Trait objects for the given identity. :param trait_data_items: list of dictionaries validated by TraitSerializerFull :param persist: determines whether the traits should be persisted to db - :return: list of TraitModels + :return: list of Traits """ trait_models = [] trait_models_to_persist = [] diff --git a/api/environments/identities/serializers.py b/api/environments/identities/serializers.py index cc97265a0af6..9aedd9ae4178 100644 --- a/api/environments/identities/serializers.py +++ b/api/environments/identities/serializers.py @@ -1,13 +1,13 @@ import typing from drf_spectacular.utils import extend_schema_field +from flagsmith_schemas.dynamodb import FeatureState as EdgeFeatureState from rest_framework import serializers from rest_framework.exceptions import ValidationError from environments.identities.models import Identity from evaluation.types import EvaluatedFeatureState from features.models import FeatureState -from util.engine_models.features.models import FeatureStateModel class IdentifierOnlyIdentitySerializer(serializers.ModelSerializer): # type: ignore[type-arg] @@ -60,6 +60,8 @@ class IdentityAllFeatureStatesMVFeatureOptionSerializer(serializers.Serializer): ) def get_value(self, instance) -> typing.Union[str, int, bool]: # type: ignore[no-untyped-def] + if isinstance(instance, typing.Mapping): + return instance["value"] # type: ignore[no-any-return] return instance.value # type: ignore[no-any-return] @@ -85,12 +87,12 @@ class IdentityAllFeatureStatesSerializer(serializers.Serializer): # type: ignor ) def get_feature_state_value( - self, instance: "EvaluatedFeatureState[FeatureState | FeatureStateModel]" + self, instance: "EvaluatedFeatureState[FeatureState | EdgeFeatureState]" ) -> typing.Union[str, int, bool]: return instance.evaluation_result["value"] # type: ignore[no-any-return] def get_overridden_by( - self, instance: "EvaluatedFeatureState[FeatureState | FeatureStateModel]" + self, instance: "EvaluatedFeatureState[FeatureState | EdgeFeatureState]" ) -> typing.Optional[str]: feature_state = instance.feature_state if not isinstance(feature_state, FeatureState): @@ -105,7 +107,7 @@ def get_overridden_by( @extend_schema_field(IdentityAllFeatureStatesSegmentSerializer) def get_segment( - self, instance: "EvaluatedFeatureState[FeatureState | FeatureStateModel]" + self, instance: "EvaluatedFeatureState[FeatureState | EdgeFeatureState]" ) -> typing.Optional[typing.Dict[str, typing.Any]]: feature_state = instance.feature_state if ( diff --git a/api/evaluation/mappers.py b/api/evaluation/mappers.py index 1dd7424d4ecb..2534cb1f4c26 100644 --- a/api/evaluation/mappers.py +++ b/api/evaluation/mappers.py @@ -12,6 +12,7 @@ from django.db.models import Prefetch, Q, prefetch_related_objects from flag_engine.context import types as engine_types +from flag_engine.context.mappers import map_any_value_to_context_value from flag_engine.segments.constants import IS_SET from flag_engine.segments.types import ConditionOperator, RuleType from pydantic import TypeAdapter @@ -26,6 +27,11 @@ from segments.types import SegmentEngineMetadata if TYPE_CHECKING: + from flagsmith_schemas.dynamodb import FeatureState as EdgeFeatureState + from flagsmith_schemas.dynamodb import ( + MultivariateFeatureStateValue as EdgeMultivariateFeatureStateValue, + ) + from edge_api.identities.models import EdgeIdentity from environments.identities.models import Identity from environments.identities.traits.models import Trait @@ -33,10 +39,6 @@ from features.models import FeatureState from features.multivariate.models import MultivariateFeatureStateValue from segments.models import Condition, Segment, SegmentRule - from util.engine_models.features.models import ( - FeatureStateModel, - MultivariateFeatureStateValueModel, - ) __all__ = ( @@ -280,19 +282,24 @@ def map_edge_identity_to_identity_context( environment: "Environment", ) -> "IdentityContext": """Map an edge identity, read back from DynamoDB, to an IdentityContext.""" - identity_model = edge_identity.engine_identity_model + document = edge_identity.document return { "identifier": edge_identity.identifier, "key": edge_identity.get_hash_key( environment.use_identity_composite_key_for_hashing ), "traits": { - trait.trait_key: trait.trait_value - for trait in identity_model.identity_traits - } - # System-owned traits are not user data: on a key clash, the system - # value wins. - | (identity_model.system_traits or {}), + trait_key: map_any_value_to_context_value(trait_value) + for trait_key, trait_value in ( + { + trait["trait_key"]: trait["trait_value"] + for trait in document["identity_traits"] + } + # System-owned traits are not user data: on a key clash, the + # system value wins. + | (document.get("system_traits") or {}) + ).items() + }, } @@ -376,25 +383,27 @@ def map_segment_to_segment_context( def map_engine_feature_state_to_feature_context( - feature_state: "FeatureStateModel", + feature_state: "EdgeFeatureState", *, priority: float | None = None, ) -> FeatureContext: - """Map a DynamoDB-sourced FeatureStateModel to a FeatureContext TypedDict. + """Map a DynamoDB-sourced feature state to a FeatureContext TypedDict. An edge identity's overrides are stored rather than evaluated, so they carry no bucketing salt: their own id seeds allocation, as it always has. """ feature_context: FeatureContext = { - "key": str(feature_state.django_id or feature_state.featurestate_uuid), - "name": feature_state.feature.name, - "enabled": feature_state.enabled, - "value": feature_state.feature_state_value, + "key": str( + feature_state.get("django_id") or feature_state.get("featurestate_uuid") + ), + "name": feature_state["feature"]["name"], + "enabled": feature_state["enabled"], + "value": feature_state["feature_state_value"], "metadata": FeatureEngineMetadata(edge_feature_state=feature_state), } if variants := _map_engine_mv_fs_values_to_feature_values( - feature_state.multivariate_feature_state_values + feature_state.get("multivariate_feature_state_values", []) ): feature_context["variants"] = variants @@ -405,24 +414,25 @@ def map_engine_feature_state_to_feature_context( def _map_engine_mv_fs_values_to_feature_values( - mv_fs_values: "Iterable[MultivariateFeatureStateValueModel]", + mv_fs_values: "Iterable[EdgeMultivariateFeatureStateValue]", ) -> list[engine_types.FeatureValue]: # Ordered by id as the stored models always have been, falling back to the # uuid for values that never reached the ORM. feature_values: list[engine_types.FeatureValue] = [] for index, mv_fs_value in enumerate( sorted( - mv_fs_values, key=lambda mv_value: mv_value.id or mv_value.mv_fs_value_uuid + mv_fs_values, + key=lambda mv_value: mv_value.get("id") or mv_value["mv_fs_value_uuid"], ) ): - mv_option = mv_fs_value.multivariate_feature_option + mv_option = mv_fs_value["multivariate_feature_option"] feature_value: engine_types.FeatureValue = { - "value": mv_option.value, - "weight": mv_fs_value.percentage_allocation, + "value": mv_option["value"], + "weight": float(mv_fs_value["percentage_allocation"]), "priority": index, } - if mv_option.key is not None: - feature_value["key"] = mv_option.key + if (key := mv_option.get("key")) is not None: + feature_value["key"] = key feature_values.append(feature_value) return feature_values diff --git a/api/evaluation/services.py b/api/evaluation/services.py index a11e4e1d6840..95c4d32c6dc5 100644 --- a/api/evaluation/services.py +++ b/api/evaluation/services.py @@ -19,13 +19,14 @@ ) if TYPE_CHECKING: + from flagsmith_schemas.dynamodb import FeatureState as EdgeFeatureState + from edge_api.identities.models import EdgeIdentity from environments.identities.models import Identity from environments.identities.traits.models import Trait from environments.models import Environment from features.models import FeatureState from segments.models import Segment - from util.engine_models.features.models import FeatureStateModel _IDENTITY_FREE_PROPERTY_PREFIXES = ( @@ -109,7 +110,7 @@ def get_environment_feature_states( def get_edge_identity_feature_states( edge_identity: "EdgeIdentity", -) -> "list[EvaluatedFeatureState[FeatureState | FeatureStateModel]]": +) -> "list[EvaluatedFeatureState[FeatureState | EdgeFeatureState]]": """The flags to serve an edge identity, one per feature.""" environment: "Environment" = edge_identity.environment @@ -146,7 +147,7 @@ def get_edge_identity_feature_states( def get_edge_identity_override_value( edge_identity: "EdgeIdentity", - feature_state: "FeatureStateModel", + feature_state: "EdgeFeatureState", *, environment: "Environment", ) -> Any: diff --git a/api/evaluation/types.py b/api/evaluation/types.py index 358adb176d91..49b6909fd3de 100644 --- a/api/evaluation/types.py +++ b/api/evaluation/types.py @@ -8,8 +8,9 @@ from segments.types import SegmentEngineMetadata if TYPE_CHECKING: + from flagsmith_schemas.dynamodb import FeatureState as EdgeFeatureState + from features.models import FeatureState - from util.engine_models.features.models import FeatureStateModel __all__ = ( @@ -39,7 +40,7 @@ FeatureStateT = TypeVar( "FeatureStateT", - bound="FeatureState | FeatureStateModel", + bound="FeatureState | EdgeFeatureState", default="FeatureState", ) diff --git a/api/features/types.py b/api/features/types.py index d82a1b506eb3..7475c7554919 100644 --- a/api/features/types.py +++ b/api/features/types.py @@ -3,8 +3,9 @@ from typing_extensions import NotRequired, TypedDict if TYPE_CHECKING: + from flagsmith_schemas.dynamodb import FeatureState as EdgeFeatureState + from features.models import FeatureState - from util.engine_models.features.models import FeatureStateModel class FeatureEngineMetadata(TypedDict): @@ -23,4 +24,4 @@ class FeatureEngineMetadata(TypedDict): feature_state: NotRequired["FeatureState"] #: An edge identity's own overrides are stored in DynamoDB rather than the #: ORM, so they reach evaluation as the model they were read back as. - edge_feature_state: NotRequired["FeatureStateModel"] + edge_feature_state: NotRequired["EdgeFeatureState"] diff --git a/api/metrics/metrics_service.py b/api/metrics/metrics_service.py index 2f24909d7d48..8f90191de3ec 100644 --- a/api/metrics/metrics_service.py +++ b/api/metrics/metrics_service.py @@ -100,7 +100,7 @@ def _get_active_identity_edge_overrides_count(self) -> int: count = 0 for override in all_overrides: - if override.feature_state.feature.id in environment_feature_ids: + if override["feature_state"]["feature"]["id"] in environment_feature_ids: count += 1 return count diff --git a/api/tests/conftest.py b/api/tests/conftest.py index cfc3f05a3f42..121a9babcc13 100644 --- a/api/tests/conftest.py +++ b/api/tests/conftest.py @@ -66,10 +66,6 @@ DynamoEnvironmentWrapper, DynamoIdentityWrapper, ) -from environments.dynamodb.types import IdentityOverrideV2 -from environments.dynamodb.utils import ( - get_environments_v2_identity_override_document_key, -) from environments.identities.models import Identity from environments.identities.traits.models import Trait from environments.models import Environment, EnvironmentAPIKey @@ -131,6 +127,7 @@ ) from users.models import FFAdminUser, UserPermissionGroup from util.mappers import ( + map_engine_feature_state_to_identity_override, map_environment_to_environment_document, map_environment_to_environment_v2_document, map_feature_state_to_engine, @@ -1623,14 +1620,7 @@ def identity_override_document( identity_uuid = str(uuid.uuid4()) identifier = "identity-with-dynamo-override" - identity_override = IdentityOverrideV2( - environment_id=str(environment.id), - environment_api_key=environment.api_key, - document_key=get_environments_v2_identity_override_document_key( - feature_id=feature.id, identity_uuid=identity_uuid - ), - identifier=identifier, - identity_uuid=identity_uuid, + identity_override = map_engine_feature_state_to_identity_override( feature_state=map_feature_state_to_engine( FeatureState( feature=feature, @@ -1638,6 +1628,10 @@ def identity_override_document( environment=environment, ), ), + identity_uuid=identity_uuid, + identifier=identifier, + environment_api_key=environment.api_key, + environment_id=environment.id, ) identity_override_document = map_identity_override_to_identity_override_document( diff --git a/api/tests/integration/edge_api/identities/conftest.py b/api/tests/integration/edge_api/identities/conftest.py index 70accd3e62d5..96443cb0c5ba 100644 --- a/api/tests/integration/edge_api/identities/conftest.py +++ b/api/tests/integration/edge_api/identities/conftest.py @@ -8,7 +8,6 @@ DynamoEnvironmentV2Wrapper, ) from users.models import FFAdminUser -from util.engine_models.identities.models import IdentityModel @pytest.fixture() @@ -27,9 +26,9 @@ def identity_overrides_v2( admin_user: FFAdminUser, ) -> list[str]: edge_identity = EdgeIdentity.from_identity_document(identity_document_without_fs) - for feature_override in IdentityModel.model_validate( + for feature_override in EdgeIdentity.from_identity_document( identity_document - ).identity_features: + ).feature_overrides: edge_identity.add_feature_override(feature_override) edge_identity.save(admin_user) return [ diff --git a/api/tests/integration/edge_api/identities/test_edge_identity_featurestates_viewset.py b/api/tests/integration/edge_api/identities/test_edge_identity_featurestates_viewset.py index b70bb1d7cbd3..76290048d37b 100644 --- a/api/tests/integration/edge_api/identities/test_edge_identity_featurestates_viewset.py +++ b/api/tests/integration/edge_api/identities/test_edge_identity_featurestates_viewset.py @@ -1,7 +1,6 @@ import copy import json import typing -import uuid from unittest import mock import pytest @@ -15,11 +14,7 @@ from rest_framework.test import APIClient from core.constants import BOOLEAN, INTEGER, STRING -from edge_api.identities.models import ( # type: ignore[attr-defined] - EdgeIdentity, - IdentityFeaturesList, - IdentityModel, -) +from edge_api.identities.models import EdgeIdentity, new_feature_override from environments.dynamodb import ( DynamoEnvironmentV2Wrapper, DynamoIdentityWrapper, @@ -29,13 +24,6 @@ from features.multivariate.models import MultivariateFeatureOption from projects.models import Project from tests.integration.helpers import create_mv_option_with_api -from util.engine_models.features.models import ( - FeatureModel, - FeatureStateModel, - MultivariateFeatureOptionModel, - MultivariateFeatureStateValueList, - MultivariateFeatureStateValueModel, -) from util.mappers.engine import map_feature_to_engine @@ -1133,13 +1121,7 @@ def test_edge_identity_clone_flag_states_from__source_with_overrides__clones_to_ ) def create_identity(identifier: str) -> EdgeIdentity: - identity_model = IdentityModel( - identifier=identifier, - environment_api_key=environment_api_key, - identity_features=IdentityFeaturesList(), - identity_uuid=uuid.uuid4(), - ) - return EdgeIdentity(engine_identity_model=identity_model) + return EdgeIdentity.create(identifier, environment_api_key) def features_for_identity_clone_flag_states_from( project: Project, @@ -1174,56 +1156,48 @@ def features_for_identity_clone_flag_states_from( string_value="bar", ) - feature_model_1: FeatureModel = map_feature_to_engine(feature=feature_1) - feature_model_2: FeatureModel = map_feature_to_engine(feature=feature_2) - feature_model_3: FeatureModel = map_feature_to_engine(feature=feature_3) - mv_feature_model: FeatureModel = map_feature_to_engine(feature=mv_feature) + feature_model_1 = map_feature_to_engine(feature=feature_1) + feature_model_2 = map_feature_to_engine(feature=feature_2) + feature_model_3 = map_feature_to_engine(feature=feature_3) + mv_feature_model = map_feature_to_engine(feature=mv_feature) source_identity: EdgeIdentity = create_identity(identifier="source_identity") target_identity: EdgeIdentity = create_identity(identifier="target_identity") source_feature_state_1_value = "Source Identity for feature value 1" - source_feature_state_1 = FeatureStateModel( # type: ignore[call-arg] + source_feature_state_1 = new_feature_override( feature=feature_model_1, - environment_id=dynamo_enabled_environment, enabled=True, feature_state_value=source_feature_state_1_value, ) source_feature_state_2_value = "Source Identity for feature value 2" - source_feature_state_2 = FeatureStateModel( # type: ignore[call-arg] + source_feature_state_2 = new_feature_override( feature=feature_model_2, - environment_id=dynamo_enabled_environment, enabled=True, feature_state_value=source_feature_state_2_value, ) - source_mv_feature_state = FeatureStateModel( # type: ignore[call-arg] + source_mv_feature_state = new_feature_override( feature=mv_feature_model, - environment_id=dynamo_enabled_environment, enabled=True, - multivariate_feature_state_values=MultivariateFeatureStateValueList(), - ) - source_mv_feature_state.multivariate_feature_state_values.append( - MultivariateFeatureStateValueModel( - multivariate_feature_option=MultivariateFeatureOptionModel( - value=mv_variant_1.value - ), - percentage_allocation=100, - ) + multivariate_feature_state_values=[ + { + "multivariate_feature_option": {"value": mv_variant_1.value}, + "percentage_allocation": 100, + } + ], ) target_feature_state_2_value = "Target Identity value for feature 2" - target_feature_state_2 = FeatureStateModel( # type: ignore[call-arg] + target_feature_state_2 = new_feature_override( feature=feature_model_2, - environment_id=dynamo_enabled_environment, enabled=False, feature_state_value=target_feature_state_2_value, ) - target_feature_state_3 = FeatureStateModel( # type: ignore[call-arg] + target_feature_state_3 = new_feature_override( feature=feature_model_3, - environment_id=dynamo_enabled_environment, enabled=False, ) @@ -1268,12 +1242,12 @@ def features_for_identity_clone_flag_states_from( assert len(response) == 4 assert response[0]["feature"]["id"] == feature_1.id - assert response[0]["enabled"] == source_feature_state_1.enabled + assert response[0]["enabled"] == source_feature_state_1["enabled"] assert response[0]["feature_state_value"] == source_feature_state_1_value assert response[0]["overridden_by"] == "IDENTITY" assert response[1]["feature"]["id"] == feature_2.id - assert response[1]["enabled"] == source_feature_state_2.enabled + assert response[1]["enabled"] == source_feature_state_2["enabled"] assert response[1]["feature_state_value"] == source_feature_state_2_value assert response[1]["overridden_by"] == "IDENTITY" @@ -1283,7 +1257,7 @@ def features_for_identity_clone_flag_states_from( assert response[2]["overridden_by"] is None assert response[3]["feature"]["id"] == mv_feature.id - assert response[3]["enabled"] == source_mv_feature_state.enabled + assert response[3]["enabled"] == source_mv_feature_state["enabled"] assert response[3]["feature_state_value"] == mv_variant_1.value assert ( response[3]["multivariate_feature_state_values"][0][ diff --git a/api/tests/unit/edge_api/identities/test_edge_api_identities_serializers.py b/api/tests/unit/edge_api/identities/test_edge_api_identities_serializers.py index c75e88fb332d..981bc601888d 100644 --- a/api/tests/unit/edge_api/identities/test_edge_api_identities_serializers.py +++ b/api/tests/unit/edge_api/identities/test_edge_api_identities_serializers.py @@ -5,7 +5,7 @@ from pytest_mock import MockerFixture from api_keys.user import APIKeyUser -from edge_api.identities.models import EdgeIdentity +from edge_api.identities.models import EdgeIdentity, new_feature_override from edge_api.identities.serializers import EdgeIdentityFeatureStateSerializer from environments.identities.models import Identity from environments.identities.serializers import ( @@ -15,7 +15,6 @@ from features.feature_types import STANDARD from features.models import Feature from users.models import FFAdminUser -from util.engine_models.features.models import FeatureModel, FeatureStateModel from util.mappers import map_identity_to_identity_document from webhooks.constants import WEBHOOK_DATETIME_FORMAT @@ -139,11 +138,11 @@ def test_edge_identity_feature_state_serializer__update_override__calls_webhook( new_enabled_state = True new_value = "bar" - instance = FeatureStateModel( - feature=FeatureModel(id=feature.id, name=feature.name, type=STANDARD), + instance = new_feature_override( + feature={"id": feature.id, "name": feature.name, "type": STANDARD}, enabled=previous_enabled_state, + feature_state_value=previous_value, ) - instance.set_value(previous_value) serializer = EdgeIdentityFeatureStateSerializer( instance=instance, diff --git a/api/tests/unit/edge_api/identities/test_edge_identity_models.py b/api/tests/unit/edge_api/identities/test_edge_identity_models.py index dc30ca9670e0..6e75c59e9524 100644 --- a/api/tests/unit/edge_api/identities/test_edge_identity_models.py +++ b/api/tests/unit/edge_api/identities/test_edge_identity_models.py @@ -13,7 +13,7 @@ from task_processor.task_run_method import TaskRunMethod from api_keys.user import APIKeyUser -from edge_api.identities.models import EdgeIdentity +from edge_api.identities.models import EdgeIdentity, new_feature_override from environments.models import Environment from evaluation.services import get_edge_identity_feature_states from features.models import Feature, FeatureSegment, FeatureState @@ -23,9 +23,6 @@ from segments.models import Condition, Segment, SegmentRule from tests.types import EnableFeaturesFixture from users.models import FFAdminUser -from util.engine_models.features.models import FeatureModel, FeatureStateModel -from util.engine_models.identities.models import IdentityModel -from util.engine_models.identities.traits.models import TraitModel MATCHING_TRAIT_KEY = "segment-membership" MATCHING_TRAIT_VALUE = "yes" @@ -43,12 +40,12 @@ def _create_matching_segment(project: Project, name: str) -> Segment: return segment -def _matching_identity_model(environment_api_key: str) -> IdentityModel: - return IdentityModel( - identifier="identity", - environment_api_key=environment_api_key, +def _matching_edge_identity(environment_api_key: str) -> EdgeIdentity: + return EdgeIdentity.create( + "identity", + environment_api_key, identity_traits=[ - TraitModel(trait_key=MATCHING_TRAIT_KEY, trait_value=MATCHING_TRAIT_VALUE) + {"trait_key": MATCHING_TRAIT_KEY, "trait_value": MATCHING_TRAIT_VALUE} ], ) @@ -76,7 +73,7 @@ def test_get_all_feature_states__multiple_segment_overrides__uses_segment_priori feature=feature, environment=environment, feature_segment=feature_segment_p2 ) - edge_identity = EdgeIdentity(_matching_identity_model(environment.api_key)) + edge_identity = _matching_edge_identity(environment.api_key) # When feature_states = get_edge_identity_feature_states(edge_identity) @@ -106,10 +103,7 @@ def test_get_all_feature_states__not_live_change_request__ignores_not_live_state change_request=change_request, ) - identity_model = mocker.MagicMock( - environment_api_key=environment.api_key, identity_features=[] - ) - edge_identity = EdgeIdentity(identity_model) + edge_identity = EdgeIdentity.create("identity", environment.api_key) # When with freeze_time(timezone.now() + timedelta(hours=2)): @@ -151,8 +145,8 @@ def test_edge_identity_id__parametrised_ids__returns_expected_id( # type: ignor django_id, identity_uuid, expected_id, mocker ): # Given / When - edge_identity = EdgeIdentity( - mocker.MagicMock(django_id=django_id, identity_uuid=identity_uuid) + edge_identity = EdgeIdentity.create( + "identity", "api-key", django_id=django_id, identity_uuid=identity_uuid ) # Then @@ -163,18 +157,18 @@ def test_get_feature_state_by_feature_name_or_id__existing_override__returns_fea edge_identity_model, ): # Given - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) edge_identity_model.add_feature_override(feature_state_model) # When found_by_name = edge_identity_model.get_feature_state_by_feature_name_or_id( - feature_state_model.feature.name + feature_state_model["feature"]["name"] ) found_by_id = edge_identity_model.get_feature_state_by_feature_name_or_id( - feature_state_model.feature.id + feature_state_model["feature"]["id"] ) # Then @@ -188,8 +182,8 @@ def test_get_feature_state_by_featurestate_uuid__existing_override__returns_feat edge_identity_model, ): # Given - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) edge_identity_model.add_feature_override(feature_state_model) @@ -197,7 +191,7 @@ def test_get_feature_state_by_featurestate_uuid__existing_override__returns_feat # When found_by_feature_state_uuid = ( edge_identity_model.get_feature_state_by_featurestate_uuid( - str(feature_state_model.featurestate_uuid) + str(feature_state_model["featurestate_uuid"]) ) ) @@ -210,8 +204,8 @@ def test_remove_feature_override__existing_override__removes_feature_state( # t edge_identity_model, ): # Given - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) edge_identity_model.add_feature_override(feature_state_model) @@ -222,7 +216,7 @@ def test_remove_feature_override__existing_override__removes_feature_state( # t # Then assert ( edge_identity_model.get_feature_state_by_feature_name_or_id( - feature_state_model.feature.id + feature_state_model["feature"]["id"] ) is None ) @@ -232,8 +226,8 @@ def test_remove_feature_override__no_matching_override__no_error( # type: ignor edge_identity_model, ): # Given - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) @@ -243,7 +237,7 @@ def test_remove_feature_override__no_matching_override__no_error( # type: ignor # Then assert ( edge_identity_model.get_feature_state_by_feature_name_or_id( - feature_state_model.feature.id + feature_state_model["feature"]["id"] ) is None ) @@ -257,8 +251,8 @@ def test_synchronise_features__empty_feature_list__removes_overrides( # type: i "edge_api.identities.models.sync_identity_document_features" ) - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) edge_identity_model.add_feature_override(feature_state_model) @@ -269,7 +263,7 @@ def test_synchronise_features__empty_feature_list__removes_overrides( # type: i # Then assert ( edge_identity_model.get_feature_state_by_feature_name_or_id( - feature_state_model.feature.id + feature_state_model["feature"]["id"] ) is None ) @@ -319,8 +313,8 @@ def test_edge_identity_save_called__feature_override_added__expected_tasks_calle "edge_api.identities.models.update_flagsmith_environments_v2_identity_overrides" ) - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) edge_identity_model.add_feature_override(feature_state_model) @@ -330,7 +324,7 @@ def test_edge_identity_save_called__feature_override_added__expected_tasks_calle "test_feature": { "change_type": "+", "new": { - **feature_state_model.dict(), + **feature_state_model, "enabled": True, "feature_state_value": None, }, @@ -388,8 +382,8 @@ def test_edge_identity_save_called__feature_override_removed__expected_tasks_cal "edge_api.identities.models.update_flagsmith_environments_v2_identity_overrides" ) - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=True, ) edge_identity_model.add_feature_override(feature_state_model) @@ -399,7 +393,7 @@ def test_edge_identity_save_called__feature_override_removed__expected_tasks_cal "test_feature": { "change_type": "-", "old": { - **feature_state_model.dict(), + **feature_state_model, "enabled": True, "feature_state_value": None, }, @@ -466,11 +460,11 @@ def test_save__feature_override_updated__generates_audit_records( "edge_api.identities.models.update_flagsmith_environments_v2_identity_overrides" ) - feature_state_model = FeatureStateModel( - feature=FeatureModel(id=1, name="test_feature", type="STANDARD"), + feature_state_model = new_feature_override( + feature={"id": 1, "name": "test_feature", "type": "STANDARD"}, enabled=initial_enabled, ) - feature_state_model.set_value(initial_value) + feature_state_model["feature_state_value"] = initial_value edge_identity_model.add_feature_override(feature_state_model) user = mocker.MagicMock() @@ -480,12 +474,12 @@ def test_save__feature_override_updated__generates_audit_records( "test_feature": { "change_type": "~", "old": { - **feature_state_model.dict(), + **feature_state_model, "enabled": initial_enabled, "feature_state_value": initial_value, }, "new": { - **feature_state_model.dict(), + **feature_state_model, "enabled": new_enabled, "feature_state_value": new_value, }, @@ -500,10 +494,11 @@ def test_save__feature_override_updated__generates_audit_records( mocked_update_flagsmith_environments_v2_identity_overrides.reset_mock() feature_override = edge_identity_model.get_feature_state_by_featurestate_uuid( - str(feature_state_model.featurestate_uuid) + str(feature_state_model["featurestate_uuid"]) ) - feature_override.enabled = new_enabled # type: ignore[union-attr] - feature_override.set_value(new_value) # type: ignore[union-attr] + assert feature_override + feature_override["enabled"] = new_enabled + feature_override["feature_state_value"] = new_value # When edge_identity_model.save(user=admin_user) @@ -561,7 +556,7 @@ def test_get_all_feature_states__post_v2_versioning_migration__returns_latest_ov operator=EQUAL, value=MATCHING_TRAIT_VALUE, ) - edge_identity = EdgeIdentity(_matching_identity_model(environment.api_key)) + edge_identity = _matching_edge_identity(environment.api_key) # When with django_assert_num_queries(8): diff --git a/api/tests/unit/edge_api/identities/test_unit_edge_api_identities_tasks.py b/api/tests/unit/edge_api/identities/test_unit_edge_api_identities_tasks.py index 4b936d3c16ea..596fb21a9811 100644 --- a/api/tests/unit/edge_api/identities/test_unit_edge_api_identities_tasks.py +++ b/api/tests/unit/edge_api/identities/test_unit_edge_api_identities_tasks.py @@ -15,12 +15,12 @@ ) from environments.dynamodb.types import ( IdentityOverridesV2Changeset, - IdentityOverrideV2, ) from environments.identities.models import Identity from environments.models import Environment, Webhook from features.models import Feature from users.models import FFAdminUser +from util.mappers import map_identity_override_document_to_identity_override from webhooks.webhooks import WebhookEventType @@ -451,13 +451,14 @@ def test_update_flagsmith_environments_v2_identity_overrides__changes_provided__ } expected_identity_overrides_changeset = IdentityOverridesV2Changeset( to_delete=[ - IdentityOverrideV2.parse_obj( + map_identity_override_document_to_identity_override( { "document_key": f"identity_override:3:{identity_uuid}", "environment_id": str(environment.id), "environment_api_key": environment.api_key, "identifier": identifier, "identity_uuid": identity_uuid, + "created_date": timezone.now(), "feature_state": { "enabled": True, "feature_state_value": "deleted", @@ -472,13 +473,14 @@ def test_update_flagsmith_environments_v2_identity_overrides__changes_provided__ ) ], to_put=[ - IdentityOverrideV2.parse_obj( + map_identity_override_document_to_identity_override( { "document_key": f"identity_override:1:{identity_uuid}", "environment_id": str(environment.id), "environment_api_key": environment.api_key, "identifier": identifier, "identity_uuid": identity_uuid, + "created_date": timezone.now(), "feature_state": { "enabled": True, "feature_state_value": "updated", @@ -491,13 +493,14 @@ def test_update_flagsmith_environments_v2_identity_overrides__changes_provided__ }, } ), - IdentityOverrideV2.parse_obj( + map_identity_override_document_to_identity_override( { "document_key": f"identity_override:2:{identity_uuid}", "environment_id": str(environment.id), "environment_api_key": environment.api_key, "identifier": identifier, "identity_uuid": identity_uuid, + "created_date": timezone.now(), "feature_state": { "enabled": True, "feature_state_value": "new", diff --git a/api/tests/unit/environments/dynamodb/test_unit_services.py b/api/tests/unit/environments/dynamodb/test_unit_services.py index 1d7478f43380..196dd6b23a7a 100644 --- a/api/tests/unit/environments/dynamodb/test_unit_services.py +++ b/api/tests/unit/environments/dynamodb/test_unit_services.py @@ -46,9 +46,9 @@ def test_migrate_environments_to_v2__environment_with_overrides__writes_expected expected_identity_override_document = ( map_identity_override_to_identity_override_document( map_engine_feature_state_to_identity_override( - feature_state=engine_identity.identity_features[0], - identity_uuid=str(engine_identity.identity_uuid), - identifier=engine_identity.identifier, + feature_state=engine_identity["identity_features"][0], + identity_uuid=str(engine_identity["identity_uuid"]), + identifier=engine_identity["identifier"], environment_api_key=environment.api_key, environment_id=environment.id, ), diff --git a/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py index bbf9cbff52aa..1cae33e27830 100644 --- a/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py +++ b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_environment_v2_wrapper.py @@ -12,7 +12,6 @@ from environments.dynamodb import DynamoEnvironmentV2Wrapper from environments.dynamodb.types import ( IdentityOverridesV2Changeset, - IdentityOverrideV2, ) from environments.dynamodb.utils import ( get_environments_v2_identity_override_document_key, @@ -21,8 +20,10 @@ from features.models import Feature, FeatureState from tests.types import EnableFeaturesFixture from util.mappers import ( + map_engine_feature_state_to_identity_override, map_environment_to_environment_v2_document, map_feature_state_to_engine, + map_identity_override_document_to_identity_override, map_identity_override_to_identity_override_document, ) @@ -96,17 +97,12 @@ def test_environment_v2_wrapper__update_identity_overrides__put_expected( identity_uuid = str(uuid.uuid4()) identifier = "identity1" - override_document = IdentityOverrideV2.parse_obj( - { - "environment_id": str(environment.id), - "document_key": get_environments_v2_identity_override_document_key( - feature_id=feature.id, identity_uuid=identity_uuid - ), - "environment_api_key": environment.api_key, - "feature_state": map_feature_state_to_engine(feature_state), - "identifier": identifier, - "identity_uuid": identity_uuid, - } + override_document = map_engine_feature_state_to_identity_override( + feature_state=map_feature_state_to_engine(feature_state), + identity_uuid=identity_uuid, + identifier=identifier, + environment_api_key=environment.api_key, + environment_id=environment.id, ) # When @@ -138,17 +134,12 @@ def test_environment_v2_wrapper__update_identity_overrides_put_frozen_time__stor wrapper = DynamoEnvironmentV2Wrapper() identity_uuid = str(uuid.uuid4()) - override_document = IdentityOverrideV2.parse_obj( - { - "environment_id": str(environment.id), - "document_key": get_environments_v2_identity_override_document_key( - feature_id=feature.id, identity_uuid=identity_uuid - ), - "environment_api_key": environment.api_key, - "feature_state": map_feature_state_to_engine(feature_state), - "identifier": "identity1", - "identity_uuid": identity_uuid, - } + override_document = map_engine_feature_state_to_identity_override( + feature_state=map_feature_state_to_engine(feature_state), + identity_uuid=identity_uuid, + identifier="identity1", + environment_api_key=environment.api_key, + environment_id=environment.id, ) # When @@ -182,23 +173,20 @@ def test_environment_v2_wrapper__update_identity_overrides__delete_expected( identity_uuid = str(uuid.uuid4()) identifier = "identity1" override_document_data = map_identity_override_to_identity_override_document( - IdentityOverrideV2.parse_obj( - { - "environment_id": str(environment.id), - "document_key": get_environments_v2_identity_override_document_key( - feature_id=feature.id, identity_uuid=identity_uuid - ), - "environment_api_key": environment.api_key, - "feature_state": map_feature_state_to_engine(feature_state), - "identifier": identifier, - "identity_uuid": identity_uuid, - } + map_engine_feature_state_to_identity_override( + feature_state=map_feature_state_to_engine(feature_state), + identity_uuid=identity_uuid, + identifier=identifier, + environment_api_key=environment.api_key, + environment_id=environment.id, ) ) flagsmith_environments_v2_table.put_item(Item=override_document_data) - override_document = IdentityOverrideV2.parse_obj(override_document_data) + override_document = map_identity_override_document_to_identity_override( + override_document_data + ) # When wrapper.update_identity_overrides( diff --git a/api/tests/unit/evaluation/test_unit_evaluation_mappers.py b/api/tests/unit/evaluation/test_unit_evaluation_mappers.py index aca4da998323..ab549f406367 100644 --- a/api/tests/unit/evaluation/test_unit_evaluation_mappers.py +++ b/api/tests/unit/evaluation/test_unit_evaluation_mappers.py @@ -3,7 +3,7 @@ from pytest_django import DjangoAssertNumQueries from pytest_lazy_fixtures import lf as lazy_fixture -from edge_api.identities.models import EdgeIdentity +from edge_api.identities.models import EdgeIdentity, new_feature_override from environments.identities.models import Identity from environments.identities.traits.models import Trait from environments.models import Environment @@ -18,15 +18,6 @@ from features.multivariate.models import MultivariateFeatureStateValue from projects.models import Project from segments.models import Condition, Segment, SegmentRule -from util.engine_models.features.models import ( - FeatureModel, - FeatureStateModel, - MultivariateFeatureOptionModel, - MultivariateFeatureStateValueList, - MultivariateFeatureStateValueModel, -) -from util.engine_models.identities.models import IdentityModel -from util.engine_models.identities.traits.models import TraitModel def test_map_environment_to_evaluation_context__full_environment__returns_expected_context( @@ -472,16 +463,14 @@ def test_map_edge_identity_to_identity_context__system_traits__merged_with_syste environment: Environment, ) -> None: # Given - edge_identity = EdgeIdentity( - IdentityModel( - identifier="identity", - environment_api_key=environment.api_key, - identity_traits=[ - TraitModel(trait_key="owned-by-user", trait_value="user value"), - TraitModel(trait_key="clashing", trait_value="user value"), - ], - system_traits={"clashing": "system value"}, - ) + edge_identity = EdgeIdentity.create( + "identity", + environment.api_key, + identity_traits=[ + {"trait_key": "owned-by-user", "trait_value": "user value"}, + {"trait_key": "clashing", "trait_value": "user value"}, + ], + system_traits={"clashing": "system value"}, ) # When @@ -501,29 +490,27 @@ def test_map_engine_feature_state_to_feature_context__multivariate_override__ret None ): # Given - feature_state = FeatureStateModel( + feature_state = new_feature_override( django_id=1, - feature=FeatureModel(id=1, name="feature", type="MULTIVARIATE"), + feature={"id": 1, "name": "feature", "type": "MULTIVARIATE"}, enabled=True, feature_state_value="control", - multivariate_feature_state_values=MultivariateFeatureStateValueList( - [ - MultivariateFeatureStateValueModel( - id=3, - percentage_allocation=70, - multivariate_feature_option=MultivariateFeatureOptionModel( - id=3, value="unkeyed" - ), - ), - MultivariateFeatureStateValueModel( - id=2, - percentage_allocation=30, - multivariate_feature_option=MultivariateFeatureOptionModel( - id=2, value="keyed", key="variant-a" - ), - ), - ] - ), + multivariate_feature_state_values=[ + { + "id": 3, + "percentage_allocation": 70, + "multivariate_feature_option": {"id": 3, "value": "unkeyed"}, + }, + { + "id": 2, + "percentage_allocation": 30, + "multivariate_feature_option": { + "id": 2, + "value": "keyed", + "key": "variant-a", + }, + }, + ], ) # When diff --git a/api/tests/unit/evaluation/test_unit_evaluation_services.py b/api/tests/unit/evaluation/test_unit_evaluation_services.py index ecbf8afef5f8..c5fb5d1d4109 100644 --- a/api/tests/unit/evaluation/test_unit_evaluation_services.py +++ b/api/tests/unit/evaluation/test_unit_evaluation_services.py @@ -2,7 +2,7 @@ from flag_engine.segments.constants import EQUAL, IN, IS_SET from pytest_lazy_fixtures import lf as lazy_fixture -from edge_api.identities.models import EdgeIdentity +from edge_api.identities.models import EdgeIdentity, new_feature_override from environments.identities.models import Identity from environments.identities.traits.models import Trait from environments.models import Environment @@ -22,15 +22,6 @@ from features.value_types import INTEGER, STRING from projects.models import Project from segments.models import Condition, Segment, SegmentRule -from util.engine_models.features.models import ( - FeatureModel, - FeatureStateModel, - MultivariateFeatureOptionModel, - MultivariateFeatureStateValueList, - MultivariateFeatureStateValueModel, -) -from util.engine_models.identities.models import IdentityFeaturesList, IdentityModel -from util.engine_models.identities.traits.models import TraitModel from util.mappers import map_identity_to_identity_document @@ -312,26 +303,20 @@ def test_get_edge_identity_feature_states__segment_and_identity_override__identi segment_override.feature_state_value.save() # and an identity override, stored against the identity in DynamoDB - edge_identity = EdgeIdentity( - IdentityModel( - identifier="identity", - environment_api_key=environment.api_key, - identity_traits=[ - TraitModel(trait_key=trait.trait_key, trait_value=trait.trait_value) - ], - identity_features=IdentityFeaturesList( - [ - FeatureStateModel( - django_id=1, - feature=FeatureModel( - id=feature.id, name=feature.name, type=feature.type - ), - enabled=True, - feature_state_value="identity", - ) - ] - ), - ) + edge_identity = EdgeIdentity.create( + "identity", + environment.api_key, + identity_traits=[ + {"trait_key": trait.trait_key, "trait_value": trait.trait_value} + ], + identity_features=[ + new_feature_override( + django_id=1, + feature={"id": feature.id, "name": feature.name, "type": feature.type}, + enabled=True, + feature_state_value="identity", + ) + ], ) # When @@ -504,35 +489,27 @@ def test_get_edge_identity_override_value__override__returns_served_value( ) -> None: # Given options = multivariate_feature.multivariate_options.order_by("id") - override = FeatureStateModel( - feature=FeatureModel( - id=multivariate_feature.id, - name=multivariate_feature.name, - type=multivariate_feature.type, - ), + override = new_feature_override( + feature={ + "id": multivariate_feature.id, + "name": multivariate_feature.name, + "type": multivariate_feature.type, + }, enabled=True, feature_state_value="control", - multivariate_feature_state_values=MultivariateFeatureStateValueList( - [ - MultivariateFeatureStateValueModel( - id=id_, - percentage_allocation=allocation, - multivariate_feature_option=MultivariateFeatureOptionModel( - id=option.id, value=option.value - ), - ) - for id_, (option, allocation) in enumerate( - zip(options, allocations), start=1 - ) - ] - ), + multivariate_feature_state_values=[ + { + "id": id_, + "percentage_allocation": allocation, + "multivariate_feature_option": {"id": option.id, "value": option.value}, + } + for id_, (option, allocation) in enumerate( + zip(options, allocations), start=1 + ) + ], ) - edge_identity = EdgeIdentity( - IdentityModel( - identifier="identity", - environment_api_key=environment.api_key, - identity_features=IdentityFeaturesList([override]), - ) + edge_identity = EdgeIdentity.create( + "identity", environment.api_key, identity_features=[override] ) (served,) = [ evaluated_feature_state diff --git a/api/tests/unit/features/test_unit_features_features_service.py b/api/tests/unit/features/test_unit_features_features_service.py index a046f2936717..7a0a7552cdee 100644 --- a/api/tests/unit/features/test_unit_features_features_service.py +++ b/api/tests/unit/features/test_unit_features_features_service.py @@ -3,7 +3,7 @@ import pytest -from edge_api.identities.models import EdgeIdentity +from edge_api.identities.models import EdgeIdentity, new_feature_override from environments.identities.models import Identity from features.features_service import ( get_core_overrides_data, @@ -13,7 +13,6 @@ from features.models import Feature, FeatureSegment, FeatureState from projects.models import EdgeV2MigrationStatus from users.models import FFAdminUser -from util.engine_models.features.models import FeatureStateModel from util.mappers.engine import ( map_feature_state_to_engine, map_identity_to_engine, @@ -246,15 +245,15 @@ def test_get_edge_overrides_data__multiple_overrides__returns_correct_counts( ) -> None: # Given # replicate identity to Edge - edge_identity = EdgeIdentity(map_identity_to_engine(identity, with_overrides=False)) + edge_identity = EdgeIdentity.from_identity_document( + map_identity_to_engine(identity, with_overrides=False) + ) edge_identity.add_feature_override( - FeatureStateModel.model_validate( - map_feature_state_to_engine(identity_featurestate) - ), + new_feature_override(**map_feature_state_to_engine(identity_featurestate)), ) edge_identity.add_feature_override( - FeatureStateModel.model_validate( - map_feature_state_to_engine(distinct_identity_featurestate) + new_feature_override( + **map_feature_state_to_engine(distinct_identity_featurestate) ), ) edge_identity.save(admin_user) @@ -305,16 +304,16 @@ def test_get_edge_overrides_data__deleted_feature__skips_deleted( # type: ignor ): # Given # replicate identity to Edge - edge_identity = EdgeIdentity(map_identity_to_engine(identity, with_overrides=False)) + edge_identity = EdgeIdentity.from_identity_document( + map_identity_to_engine(identity, with_overrides=False) + ) # Create identity override for two different features edge_identity.add_feature_override( - FeatureStateModel.model_validate( - map_feature_state_to_engine(identity_featurestate) - ), + new_feature_override(**map_feature_state_to_engine(identity_featurestate)), ) edge_identity.add_feature_override( - FeatureStateModel.model_validate( - map_feature_state_to_engine(distinct_identity_featurestate) + new_feature_override( + **map_feature_state_to_engine(distinct_identity_featurestate) ), ) edge_identity.save(admin_user) diff --git a/api/tests/unit/features/test_unit_features_views.py b/api/tests/unit/features/test_unit_features_views.py index 48af2b1fca08..3e4e7054778c 100644 --- a/api/tests/unit/features/test_unit_features_views.py +++ b/api/tests/unit/features/test_unit_features_views.py @@ -4863,9 +4863,9 @@ def test_delete_feature__dynamo_identity_overrides__deletes_overrides( flagsmith_environments_v2_table.put_item( Item=map_identity_override_to_identity_override_document( map_engine_feature_state_to_identity_override( - feature_state=engine_identity.identity_features[0], - identity_uuid=str(engine_identity.identity_uuid), - identifier=engine_identity.identifier, + feature_state=engine_identity["identity_features"][0], + identity_uuid=str(engine_identity["identity_uuid"]), + identifier=engine_identity["identifier"], environment_api_key=environment.api_key, environment_id=environment.id, ), diff --git a/api/tests/unit/metrics/test_unit_metrics_service.py b/api/tests/unit/metrics/test_unit_metrics_service.py index c237d6ab20be..71c888aba175 100644 --- a/api/tests/unit/metrics/test_unit_metrics_service.py +++ b/api/tests/unit/metrics/test_unit_metrics_service.py @@ -97,9 +97,7 @@ def test_environment_metrics_service__dynamo_enabled_or_not__uses_correct_identi lambda: MagicMock(count=identity_count_mock), ) - mock_override = MagicMock() - mock_override.feature_state.feature.id = feature.id - mock_overrides = [mock_override] * 99 + mock_overrides = [{"feature_state": {"feature": {"id": feature.id}}}] * 99 dynamo_mock = MagicMock(return_value=mock_overrides) monkeypatch.setattr( diff --git a/api/tests/unit/util/engine_models/identities/test_unit_identities_models.py b/api/tests/unit/util/engine_models/identities/test_unit_identities_models.py deleted file mode 100644 index 5d26571c5c9c..000000000000 --- a/api/tests/unit/util/engine_models/identities/test_unit_identities_models.py +++ /dev/null @@ -1,31 +0,0 @@ -import pytest - -from util.engine_models.identities.models import IdentityModel - - -@pytest.mark.parametrize( - "django_id, use_identity_composite_key_for_hashing, expected_hash_key", - [ - (None, True, "api-key_identifier"), - (1, True, "api-key_identifier"), - (1, False, "1"), - (None, False, "identifier"), - ], -) -def test_identity_model_get_hash_key__django_id_and_hashing_setting__returns_expected_key( - django_id: int | None, - use_identity_composite_key_for_hashing: bool, - expected_hash_key: str, -) -> None: - # Given - identity_model = IdentityModel( - identifier="identifier", - environment_api_key="api-key", - django_id=django_id, - ) - - # When - hash_key = identity_model.get_hash_key(use_identity_composite_key_for_hashing) - - # Then - assert hash_key == expected_hash_key diff --git a/api/tests/unit/util/engine_models/identities/traits/test_unit_traits_types.py b/api/tests/unit/util/engine_models/identities/traits/test_unit_traits_types.py deleted file mode 100644 index e18cc41e44f0..000000000000 --- a/api/tests/unit/util/engine_models/identities/traits/test_unit_traits_types.py +++ /dev/null @@ -1,31 +0,0 @@ -import pytest - -from util.engine_models.identities.traits.types import ( - map_any_value_to_trait_value, -) - - -@pytest.mark.parametrize( - "value, expected", - [ - # String values that look like integers should be converted to int - ("123", 123), - ("-45", -45), - ("0", 0), - # String values that look like floats should be converted to float - ("1.23", 1.23), - ("-4.56", -4.56), - ("0.0", 0.0), - # Non-trait-value types should be converted to string - (["a", "list"], "['a', 'list']"), - ({"a": "dict"}, "{'a': 'dict'}"), - ], -) -def test_map_any_value_to_trait_value__various_types__returns_expected_value( - value: object, expected: object -) -> None: - # Given / When - result = map_any_value_to_trait_value(value) - - # Then - assert result == expected diff --git a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py index 9351ec3ce76d..b6a328026a3d 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py @@ -2,17 +2,15 @@ import json import uuid from decimal import Decimal -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from django.utils import timezone from environments.dynamodb.constants import ( ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, ) -from util.engine_models.features.models import FeatureStateModel -from util.engine_models.identities.models import IdentityModel from util.mappers import dynamodb -from util.mappers.engine import map_feature_state_to_engine +from util.mappers.engine import map_feature_state_to_engine, map_identifier_to_engine if TYPE_CHECKING: from pytest_mock import MockerFixture @@ -160,10 +158,8 @@ def test_map_engine_identity_to_identity_document__system_traits_set__included_i None ): # Given - engine_identity = IdentityModel( - identifier="test_identity", - environment_api_key="api-key", - system_traits={"flagsmith_cohort_2b6d1f5f": True}, + engine_identity = map_identifier_to_engine( + "test_identity", "api-key", system_traits={"flagsmith_cohort_2b6d1f5f": True} ) # When @@ -177,10 +173,7 @@ def test_map_engine_identity_to_identity_document__no_system_traits__key_absent( None ): # Given - engine_identity = IdentityModel( - identifier="test_identity", - environment_api_key="api-key", - ) + engine_identity = map_identifier_to_engine("test_identity", "api-key") # When result = dynamodb.map_engine_identity_to_identity_document(engine_identity) @@ -194,18 +187,66 @@ def test_identity_document__system_traits_set__round_trip_preserves_system_trait ): # Given document = dynamodb.map_engine_identity_to_identity_document( - IdentityModel( - identifier="test_identity", - environment_api_key="api-key", + map_identifier_to_engine( + "test_identity", + "api-key", system_traits={"flagsmith_cohort_2b6d1f5f": True}, ) ) # When - parsed = IdentityModel.model_validate(document) + parsed = dynamodb.map_identity_document_to_engine_identity(document) # Then - assert parsed.system_traits == {"flagsmith_cohort_2b6d1f5f": True} + assert parsed["system_traits"] == {"flagsmith_cohort_2b6d1f5f": True} + + +def test_map_engine_identity_to_identity_document__stored_numbers__round_trip_unchanged() -> ( + None +): + # Given + stored_document = dynamodb.map_engine_identity_to_identity_document( + map_identifier_to_engine( + "test_identity", + "api-key", + identity_traits=[ + {"trait_key": "integer", "trait_value": 1}, + {"trait_key": "float", "trait_value": 1.5}, + ], + identity_features=[ + { + "feature": {"id": 1, "name": "feature", "type": "MULTIVARIATE"}, + "enabled": True, + "feature_state_value": 5, + "featurestate_uuid": str(uuid.uuid4()), + "multivariate_feature_state_values": [ + { + "mv_fs_value_uuid": str(uuid.uuid4()), + "percentage_allocation": 100, + "multivariate_feature_option": {"id": 2, "value": 3}, + } + ], + } + ], + ) + ) + + # When + document: dict[str, Any] = dynamodb.map_engine_identity_to_identity_document( + dynamodb.map_identity_document_to_engine_identity(stored_document) + ) + + # Then + assert document == stored_document + assert document["identity_traits"] == [ + {"trait_key": "integer", "trait_value": Decimal("1")}, + {"trait_key": "float", "trait_value": Decimal("1.5")}, + ] + (feature_state,) = document["identity_features"] + assert feature_state["feature_state_value"] == Decimal("5") + assert feature_state["multivariate_feature_state_values"][0][ + "multivariate_feature_option" + ]["value"] == Decimal("3") def test_map_environment_to_environment_v2_document__valid_environment__returns_expected_document( @@ -274,19 +315,17 @@ def test_map_environment_to_environment_v2_document__valid_environment__returns_ } -def test_map_identity_override_to_identity_override_document__decimal_feature_state_value__return_expected( +def test_map_identity_override_to_identity_override_document__decimal_feature_state_value__returns_string_value( identity: "Identity", identity_featurestate: "FeatureState", ) -> None: # Given - expected_feature_state_value = Decimal("1.111") + feature_state_value = Decimal("1.111") - engine_feature_state = FeatureStateModel.model_validate( - { - **map_feature_state_to_engine(identity_featurestate), - "feature_state_value": expected_feature_state_value, - } - ) + engine_feature_state = { + **map_feature_state_to_engine(identity_featurestate), + "feature_state_value": feature_state_value, + } identity_override = dynamodb.map_engine_feature_state_to_identity_override( feature_state=engine_feature_state, identity_uuid=str(uuid.uuid4()), @@ -303,9 +342,7 @@ def test_map_identity_override_to_identity_override_document__decimal_feature_st # Then feature_state = result["feature_state"] assert isinstance(feature_state, dict) - feature_state_value = feature_state["feature_state_value"] - assert isinstance(feature_state_value, Decimal) - assert feature_state_value == expected_feature_state_value + assert feature_state["feature_state_value"] == "1.111" def test_map_environment_to_compressed_environment_document__valid_environment__returns_compressed_fields( diff --git a/api/tests/unit/util/mappers/test_unit_mappers_engine.py b/api/tests/unit/util/mappers/test_unit_mappers_engine.py index 67af2beb8976..383e20dcf557 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_engine.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_engine.py @@ -20,16 +20,6 @@ from integrations.webhook.models import WebhookConfiguration from segments.models import Segment, SegmentRule from users.models import FFAdminUser -from util.engine_models.features.models import ( - FeatureModel, - FeatureStateModel, - MultivariateFeatureOptionModel, -) -from util.engine_models.identities.models import ( - IdentityFeaturesList, - IdentityModel, -) -from util.engine_models.identities.traits.models import TraitModel from util.mappers import engine if TYPE_CHECKING: @@ -314,19 +304,6 @@ def test_map_feature_state_to_engine__mv_option_with_key__key_in_serialised_docu ) -def test_multivariate_feature_option_model__document_without_key__defaults_to_none() -> ( - None -): - # Given - a document produced before `key` existed on the model - document = {"value": "control", "id": 1} - - # When - model = MultivariateFeatureOptionModel.model_validate(document) - - # Then - assert model.key is None - - def test_map_environment_to_engine__multiple_segments_and_versions__returns_expected_model( environment: Environment, feature: "Feature", @@ -564,37 +541,34 @@ def test_map_identity_to_engine__identity_with_traits_and_overrides__returns_exp ) -> None: # Given environment_api_key = environment.api_key - expected_result = IdentityModel.construct( - identifier=identity.identifier, - environment_api_key=environment_api_key, - created_date=identity.created_date, - identity_features=IdentityFeaturesList( - [ - FeatureStateModel( - feature=FeatureModel( - id=feature.pk, - name=feature.name, - type=feature.type, - ), - enabled=identity_featurestate.enabled, - django_id=identity_featurestate.pk, - feature_segment=identity_featurestate.feature_segment, - featurestate_uuid=identity_featurestate.uuid, - feature_state_value=identity_featurestate.get_feature_state_value(), - multivariate_feature_state_values=[], # type: ignore[arg-type] - ) - ] - ), - identity_traits=[ - TraitModel( - trait_key=trait.trait_key, - trait_value=trait.trait_value, - ) + expected_result = { + "identifier": identity.identifier, + "environment_api_key": environment_api_key, + "created_date": identity.created_date, + "identity_features": [ + { + "feature": { + "id": feature.pk, + "name": feature.name, + "type": feature.type, + }, + "enabled": identity_featurestate.enabled, + "django_id": identity_featurestate.pk, + "feature_segment": None, + "featurestate_uuid": identity_featurestate.uuid, + "feature_state_value": identity_featurestate.get_feature_state_value(), + "multivariate_feature_state_values": [], + } ], - identity_uuid=mocker.ANY, - django_id=identity.pk, - composite_key=identity.composite_key, - ) + "identity_traits": [ + {"trait_key": trait.trait_key, "trait_value": trait.trait_value} + ], + "system_traits": None, + "identity_uuid": mocker.ANY, + "django_id": identity.pk, + "dashboard_alias": None, + "composite_key": identity.composite_key, + } # When result = engine.map_identity_to_engine(identity) diff --git a/api/tests/unit/util/mappers/test_unit_mappers_sdk.py b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py index aa2af8f3a0d6..3a2634ecd6ee 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_sdk.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py @@ -187,7 +187,7 @@ def test_map_environment_to_sdk_document__system_traits_set__excluded_from_docum ) -> None: # Given engine_identity = map_identity_to_engine(identity, with_traits=False) - engine_identity.system_traits = {"flagsmith_cohort_2b6d1f5f": True} + engine_identity["system_traits"] = {"flagsmith_cohort_2b6d1f5f": True} mocker.patch( "util.mappers.sdk.map_identity_to_engine", return_value=engine_identity, @@ -198,5 +198,5 @@ def test_map_environment_to_sdk_document__system_traits_set__excluded_from_docum # Then assert result["identity_overrides"] == [ - engine_identity.model_dump(exclude={"system_traits"}) + {key: value for key, value in engine_identity.items() if key != "system_traits"} ] diff --git a/api/util/engine_models/__init__.py b/api/util/engine_models/__init__.py deleted file mode 100644 index a1227dbe1141..000000000000 --- a/api/util/engine_models/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -""" -Vendored Pydantic models from flagsmith-flag-engine's fix/missing-export branch. - -TEMPORARY: This module is a temporary measure to maintain compatibility during -the migration to flag-engine v10. These Pydantic models will be removed once -the codebase is fully migrated to use the TypedDict-based evaluation API -provided by flag-engine v10. -""" diff --git a/api/util/engine_models/environments/__init__.py b/api/util/engine_models/environments/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/environments/integrations/__init__.py b/api/util/engine_models/environments/integrations/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/environments/integrations/models.py b/api/util/engine_models/environments/integrations/models.py deleted file mode 100644 index aa0b456b190c..000000000000 --- a/api/util/engine_models/environments/integrations/models.py +++ /dev/null @@ -1,9 +0,0 @@ -from typing import Optional - -from pydantic import BaseModel - - -class IntegrationModel(BaseModel): - api_key: Optional[str] = None - base_url: Optional[str] = None - entity_selector: Optional[str] = None diff --git a/api/util/engine_models/environments/models.py b/api/util/engine_models/environments/models.py deleted file mode 100644 index 4198a1b3814c..000000000000 --- a/api/util/engine_models/environments/models.py +++ /dev/null @@ -1,51 +0,0 @@ -import typing -from datetime import datetime - -from pydantic import BaseModel, Field - -from util.engine_models.environments.integrations.models import IntegrationModel -from util.engine_models.features.models import FeatureStateModel -from util.engine_models.identities.models import IdentityModel -from util.engine_models.projects.models import ProjectModel -from util.engine_models.utils.datetime import utcnow_with_tz - - -class EnvironmentAPIKeyModel(BaseModel): - id: int - key: str - created_at: datetime - name: str - client_api_key: str - expires_at: typing.Optional[datetime] = None - active: bool = True - - -class WebhookModel(BaseModel): - url: str - secret: str - - -class EnvironmentModel(BaseModel): - id: int - api_key: str - project: ProjectModel - feature_states: typing.List[FeatureStateModel] = Field(default_factory=list) - identity_overrides: typing.List[IdentityModel] = Field(default_factory=list) - - name: typing.Optional[str] = None - allow_client_traits: bool = True - updated_at: datetime = Field(default_factory=utcnow_with_tz) - hide_sensitive_data: bool = False - hide_disabled_flags: typing.Optional[bool] = None - use_identity_composite_key_for_hashing: bool = False - use_identity_overrides_in_local_eval: bool = False - onboarding_pending: typing.Optional[bool] = None - - amplitude_config: typing.Optional[IntegrationModel] = None - dynatrace_config: typing.Optional[IntegrationModel] = None - heap_config: typing.Optional[IntegrationModel] = None - mixpanel_config: typing.Optional[IntegrationModel] = None - rudderstack_config: typing.Optional[IntegrationModel] = None - segment_config: typing.Optional[IntegrationModel] = None - - webhook_config: typing.Optional[WebhookModel] = None diff --git a/api/util/engine_models/features/__init__.py b/api/util/engine_models/features/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/features/models.py b/api/util/engine_models/features/models.py deleted file mode 100644 index 4bc0ce8163e0..000000000000 --- a/api/util/engine_models/features/models.py +++ /dev/null @@ -1,84 +0,0 @@ -import typing -import uuid - -from annotated_types import Ge, Le -from pydantic import UUID4, BaseModel, Field, model_validator -from pydantic_collections import BaseCollectionModel # type: ignore[import-untyped] -from typing_extensions import Annotated - -from util.engine_models.utils.exceptions import InvalidPercentageAllocation - - -class FeatureModel(BaseModel): - id: int - name: str - type: str - - -class MultivariateFeatureOptionModel(BaseModel): - value: typing.Any - id: typing.Optional[int] = None - key: typing.Optional[str] = None - - -class MultivariateFeatureStateValueModel(BaseModel): - multivariate_feature_option: MultivariateFeatureOptionModel - percentage_allocation: Annotated[float, Ge(0), Le(100)] - id: typing.Optional[int] = None - mv_fs_value_uuid: UUID4 = Field(default_factory=uuid.uuid4) - - -class FeatureSegmentModel(BaseModel): - priority: typing.Optional[int] = None - - -class MultivariateFeatureStateValueList( - BaseCollectionModel[MultivariateFeatureStateValueModel] # type: ignore[misc] -): - @staticmethod - def _ensure_correct_percentage_allocations( - value: typing.List[MultivariateFeatureStateValueModel], - ) -> typing.List[MultivariateFeatureStateValueModel]: - if ( - sum( - multivariate_feature_state.percentage_allocation - for multivariate_feature_state in value - ) - > 100 - ): - raise InvalidPercentageAllocation( - "Total percentage allocation for feature must be less or equal to 100 percent" - ) - return value - - percentage_allocations_model_validator = model_validator(mode="after")( - _ensure_correct_percentage_allocations - ) - - def append( - self, - multivariate_feature_state_value: MultivariateFeatureStateValueModel, - ) -> None: - self._ensure_correct_percentage_allocations( - [*self, multivariate_feature_state_value], - ) - super().append(multivariate_feature_state_value) - - -class FeatureStateModel(BaseModel, validate_assignment=True): - feature: FeatureModel - enabled: bool - django_id: typing.Optional[int] = None - feature_segment: typing.Optional[FeatureSegmentModel] = None - featurestate_uuid: UUID4 = Field(default_factory=uuid.uuid4) - feature_state_value: typing.Any = None - multivariate_feature_state_values: MultivariateFeatureStateValueList = Field( - default_factory=MultivariateFeatureStateValueList - ) - metadata: typing.Optional[typing.Dict[str, typing.Any]] = Field( - default=None, - exclude_if=lambda value: not value, - ) - - def set_value(self, value: typing.Any) -> None: - self.feature_state_value = value diff --git a/api/util/engine_models/identities/__init__.py b/api/util/engine_models/identities/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/identities/models.py b/api/util/engine_models/identities/models.py deleted file mode 100644 index a8d35ebb034f..000000000000 --- a/api/util/engine_models/identities/models.py +++ /dev/null @@ -1,96 +0,0 @@ -import datetime -import typing -import uuid - -from pydantic import UUID4, BaseModel, Field, computed_field, model_validator -from pydantic_collections import BaseCollectionModel # type: ignore[import-untyped] - -from util.engine_models.features.models import FeatureStateModel -from util.engine_models.identities.traits.models import TraitModel -from util.engine_models.identities.traits.types import ContextValue -from util.engine_models.utils.datetime import utcnow_with_tz -from util.engine_models.utils.exceptions import DuplicateFeatureState - - -class IdentityFeaturesList(BaseCollectionModel[FeatureStateModel]): # type: ignore[misc] - @staticmethod - def _ensure_unique_feature_ids( - value: typing.Sequence[FeatureStateModel], - ) -> None: - for i, feature_state in enumerate(value, start=1): - if feature_state.feature.id in [ - feature_state.feature.id for feature_state in value[i:] - ]: - raise DuplicateFeatureState( - f"Feature state for feature id={feature_state.feature.id} already exists" - ) - - @model_validator(mode="after") - def ensure_unique_feature_ids(self) -> "IdentityFeaturesList": - self._ensure_unique_feature_ids(self.root) - return self - - def append(self, feature_state: "FeatureStateModel") -> None: - self._ensure_unique_feature_ids([*self, feature_state]) - super().append(feature_state) - - -class IdentityModel(BaseModel): - identifier: str - environment_api_key: str - created_date: datetime.datetime = Field(default_factory=utcnow_with_tz) - identity_features: IdentityFeaturesList = Field( - default_factory=IdentityFeaturesList - ) - identity_traits: typing.List[TraitModel] = Field(default_factory=list) - # System-owned (e.g. cohort membership); unreachable by SDK and admin trait writes. - system_traits: typing.Optional[typing.Dict[str, ContextValue]] = None - identity_uuid: UUID4 = Field(default_factory=uuid.uuid4) - django_id: typing.Optional[int] = None - - dashboard_alias: typing.Optional[str] = None - - @computed_field # type: ignore[prop-decorator] - @property - def composite_key(self) -> str: - return self.generate_composite_key(self.environment_api_key, self.identifier) - - @staticmethod - def generate_composite_key(env_key: str, identifier: str) -> str: - return f"{env_key}_{identifier}" - - def get_hash_key(self, use_identity_composite_key_for_hashing: bool) -> str: - if use_identity_composite_key_for_hashing: - return self.composite_key - if self.django_id is not None: - return str(self.django_id) - return self.identifier - - def update_traits( - self, traits: typing.List[TraitModel] - ) -> typing.Tuple[typing.List[TraitModel], bool]: - existing_traits = {trait.trait_key: trait for trait in self.identity_traits} - traits_changed = False - - for trait in traits: - existing_trait = existing_traits.get(trait.trait_key) - - if trait.trait_value is None and existing_trait: - existing_traits.pop(trait.trait_key) - traits_changed = True - - elif getattr(existing_trait, "trait_value", None) != trait.trait_value: - existing_traits[trait.trait_key] = trait - traits_changed = True - - self.identity_traits = list(existing_traits.values()) - return self.identity_traits, traits_changed - - def prune_features(self, valid_feature_names: typing.List[str]) -> None: - self.identity_features = IdentityFeaturesList( - [ - fs - for fs in self.identity_features - if fs.feature.name in valid_feature_names - ] - ) diff --git a/api/util/engine_models/identities/traits/__init__.py b/api/util/engine_models/identities/traits/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/identities/traits/constants.py b/api/util/engine_models/identities/traits/constants.py deleted file mode 100644 index 359ecbd4978d..000000000000 --- a/api/util/engine_models/identities/traits/constants.py +++ /dev/null @@ -1 +0,0 @@ -TRAIT_STRING_VALUE_MAX_LENGTH: int = 2000 diff --git a/api/util/engine_models/identities/traits/models.py b/api/util/engine_models/identities/traits/models.py deleted file mode 100644 index 25069d7da767..000000000000 --- a/api/util/engine_models/identities/traits/models.py +++ /dev/null @@ -1,8 +0,0 @@ -from pydantic import BaseModel, Field - -from util.engine_models.identities.traits.types import ContextValue - - -class TraitModel(BaseModel): - trait_key: str - trait_value: ContextValue = Field(...) diff --git a/api/util/engine_models/identities/traits/types.py b/api/util/engine_models/identities/traits/types.py deleted file mode 100644 index 1b9c55f6d99e..000000000000 --- a/api/util/engine_models/identities/traits/types.py +++ /dev/null @@ -1,62 +0,0 @@ -import re -from decimal import Decimal -from typing import Any, Union, get_args - -from pydantic import BeforeValidator -from pydantic.types import AllowInfNan, StrictBool, StringConstraints -from typing_extensions import Annotated, TypeGuard - -from util.engine_models.identities.traits.constants import TRAIT_STRING_VALUE_MAX_LENGTH - -_UnconstrainedContextValue = Union[None, int, float, bool, str] - - -def map_any_value_to_trait_value(value: Any) -> _UnconstrainedContextValue: - """ - Try to coerce a value of arbitrary type to a trait value type. - Union member-specific constraints, such as max string value length, are ignored here. - Replicate behaviour from marshmallow/pydantic V1 for number-like strings. - For decimals return an int in case of unset exponent. - When in doubt, return string. - - Supposed to be used as a `pydantic.BeforeValidator`. - """ - if _is_trait_value(value): - if isinstance(value, str): - return _map_string_value_to_trait_value(value) - return value - if isinstance(value, Decimal): - if value.as_tuple().exponent: - return float(str(value)) - return int(value) - return str(value) - - -_int_pattern = re.compile(r"-?[0-9]+") -_float_pattern = re.compile(r"-?[0-9]+\.[0-9]+") - - -def _map_string_value_to_trait_value(value: str) -> _UnconstrainedContextValue: - if _int_pattern.fullmatch(value): - return int(value) - if _float_pattern.fullmatch(value): - return float(value) - return value - - -def _is_trait_value(value: Any) -> TypeGuard[_UnconstrainedContextValue]: - return isinstance(value, get_args(_UnconstrainedContextValue)) - - -ContextValue = Annotated[ - Union[ - None, - StrictBool, - Annotated[float, AllowInfNan(False)], - int, - Annotated[str, StringConstraints(max_length=TRAIT_STRING_VALUE_MAX_LENGTH)], - ], - BeforeValidator(map_any_value_to_trait_value), -] - -TraitValue = ContextValue diff --git a/api/util/engine_models/organisations/__init__.py b/api/util/engine_models/organisations/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/organisations/models.py b/api/util/engine_models/organisations/models.py deleted file mode 100644 index 1d1e9194aa9e..000000000000 --- a/api/util/engine_models/organisations/models.py +++ /dev/null @@ -1,9 +0,0 @@ -from pydantic import BaseModel - - -class OrganisationModel(BaseModel): - id: int - name: str - feature_analytics: bool = False - stop_serving_flags: bool = False - persist_trait_data: bool = True diff --git a/api/util/engine_models/projects/__init__.py b/api/util/engine_models/projects/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/projects/models.py b/api/util/engine_models/projects/models.py deleted file mode 100644 index 9a5ebd0ae9ee..000000000000 --- a/api/util/engine_models/projects/models.py +++ /dev/null @@ -1,16 +0,0 @@ -import typing - -from pydantic import BaseModel, Field - -from util.engine_models.organisations.models import OrganisationModel -from util.engine_models.segments.models import SegmentModel - - -class ProjectModel(BaseModel): - id: int - name: str - organisation: OrganisationModel - hide_disabled_flags: bool = False - segments: typing.List[SegmentModel] = Field(default_factory=list) - enable_realtime_updates: bool = False - server_key_only_feature_ids: typing.List[int] = Field(default_factory=list) diff --git a/api/util/engine_models/segments/__init__.py b/api/util/engine_models/segments/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/segments/models.py b/api/util/engine_models/segments/models.py deleted file mode 100644 index 015ca8846e07..000000000000 --- a/api/util/engine_models/segments/models.py +++ /dev/null @@ -1,28 +0,0 @@ -import typing - -from flag_engine.segments.types import ConditionOperator, RuleType -from pydantic import BaseModel, BeforeValidator, Field -from typing_extensions import Annotated - -from util.engine_models.features.models import FeatureStateModel - -LaxStr = Annotated[str, BeforeValidator(lambda x: str(x))] - - -class SegmentConditionModel(BaseModel): - operator: ConditionOperator - value: typing.Optional[LaxStr] = None - property_: typing.Optional[str] = None - - -class SegmentRuleModel(BaseModel): - type: RuleType - rules: typing.List["SegmentRuleModel"] = Field(default_factory=list) - conditions: typing.List[SegmentConditionModel] = Field(default_factory=list) - - -class SegmentModel(BaseModel): - id: int - name: str - rules: typing.List[SegmentRuleModel] = Field(default_factory=list) - feature_states: typing.List[FeatureStateModel] = Field(default_factory=list) diff --git a/api/util/engine_models/utils/__init__.py b/api/util/engine_models/utils/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/api/util/engine_models/utils/datetime.py b/api/util/engine_models/utils/datetime.py deleted file mode 100644 index 78c8697d7da1..000000000000 --- a/api/util/engine_models/utils/datetime.py +++ /dev/null @@ -1,5 +0,0 @@ -from datetime import datetime, timezone - - -def utcnow_with_tz() -> datetime: - return datetime.now(tz=timezone.utc) diff --git a/api/util/engine_models/utils/exceptions.py b/api/util/engine_models/utils/exceptions.py deleted file mode 100644 index 279e8aef095b..000000000000 --- a/api/util/engine_models/utils/exceptions.py +++ /dev/null @@ -1,6 +0,0 @@ -class DuplicateFeatureState(ValueError): - pass - - -class InvalidPercentageAllocation(ValueError): - pass diff --git a/api/util/mappers/__init__.py b/api/util/mappers/__init__.py index 7daa34fe5de7..0645b32c5510 100644 --- a/api/util/mappers/__init__.py +++ b/api/util/mappers/__init__.py @@ -7,12 +7,15 @@ map_environment_to_environment_document, map_environment_to_environment_v2_document, map_identity_changeset_to_identity_override_changeset, + map_identity_document_to_engine_identity, + map_identity_override_document_to_identity_override, map_identity_override_to_identity_override_document, map_identity_to_identity_document, ) from util.mappers.engine import ( map_feature_state_to_engine, map_feature_to_engine, + map_identifier_to_engine, map_identity_to_engine, map_mv_option_to_engine, ) @@ -29,7 +32,10 @@ "map_environment_to_sdk_document", "map_feature_state_to_engine", "map_feature_to_engine", + "map_identifier_to_engine", "map_identity_changeset_to_identity_override_changeset", + "map_identity_document_to_engine_identity", + "map_identity_override_document_to_identity_override", "map_identity_override_to_identity_override_document", "map_identity_to_engine", "map_identity_to_identity_document", diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index 7b55ea6f1464..61adecc054fd 100644 --- a/api/util/mappers/dynamodb.py +++ b/api/util/mappers/dynamodb.py @@ -1,14 +1,18 @@ -from datetime import datetime +from collections.abc import Mapping from decimal import Decimal -from typing import TYPE_CHECKING, Any, Callable, Dict, List, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast +from django.utils import timezone from flagsmith_schemas.dynamodb import ( Environment, EnvironmentAPIKey, EnvironmentCompressed, + EnvironmentV2IdentityOverride, EnvironmentV2MetaCompressed, + FeatureState, + Identity, ) -from pydantic import BaseModel, TypeAdapter +from pydantic import TypeAdapter from edge_api.identities.types import IdentityChangeset from environments.dynamodb.constants import ( @@ -23,23 +27,23 @@ get_environments_v2_identity_override_document_key, ) from util.dataclasses import CompressedEnvironmentDocument -from util.engine_models.features.models import FeatureStateModel from util.mappers.engine import ( map_environment_api_key_to_engine, map_environment_to_engine, map_identity_to_engine, ) -from util.mappers.types import Document, DocumentValue +from util.mappers.types import Document if TYPE_CHECKING: - from environments.identities.models import Identity + from environments.identities.models import Identity as IdentityModel from environments.models import Environment as EnvironmentModel from environments.models import EnvironmentAPIKey as EnvironmentAPIKeyModel - from util.engine_models.identities.models import IdentityModel __all__ = ( "map_engine_identity_to_identity_document", + "map_identity_document_to_engine_identity", + "map_identity_override_document_to_identity_override", "map_environment_api_key_to_environment_api_key_document", "map_environment_to_compressed_environment_document", "map_environment_to_compressed_environment_v2_document", @@ -49,10 +53,16 @@ ) +T = TypeVar("T") + _environment_adapter: TypeAdapter[Environment] = TypeAdapter(Environment) _environment_api_key_adapter: TypeAdapter[EnvironmentAPIKey] = TypeAdapter( EnvironmentAPIKey ) +_identity_adapter: TypeAdapter[Identity] = TypeAdapter(Identity) +_identity_override_adapter: TypeAdapter[EnvironmentV2IdentityOverride] = TypeAdapter( + EnvironmentV2IdentityOverride +) _environment_compressed_adapter: TypeAdapter[EnvironmentCompressed] = TypeAdapter( EnvironmentCompressed, ) @@ -114,42 +124,58 @@ def map_environment_api_key_to_environment_api_key_document( ) +def map_identity_document_to_engine_identity( + identity_document: Mapping[str, Any], +) -> Identity: + return _validate_document(_identity_adapter, identity_document) + + def map_engine_identity_to_identity_document( - engine_identity: "IdentityModel", + engine_identity: Mapping[str, Any], ) -> Document: - response = { - field_name: _map_value_to_document_value(value) - for field_name, value in engine_identity - if (value is not None or field_name not in _NULLABLE_IDENTITY_KEY_ATTRIBUTES) + identity_document = cast( + Document, _validate_document(_identity_adapter, engine_identity) + ) + return { + field_name: value + for field_name, value in identity_document.items() + if value is not None or field_name not in _NULLABLE_IDENTITY_KEY_ATTRIBUTES } - response["composite_key"] = engine_identity.composite_key - return response def map_identity_to_identity_document( - identity: "Identity", + identity: "IdentityModel", ) -> Document: return map_engine_identity_to_identity_document(map_identity_to_engine(identity)) +def map_identity_override_document_to_identity_override( + identity_override_document: Mapping[str, Any], +) -> IdentityOverrideV2: + return _validate_document(_identity_override_adapter, identity_override_document) + + def map_engine_feature_state_to_identity_override( *, - feature_state: "FeatureStateModel", + feature_state: Mapping[str, Any] | FeatureState, identity_uuid: str, identifier: str, environment_api_key: str, environment_id: int, ) -> IdentityOverrideV2: - return IdentityOverrideV2( - document_key=get_environments_v2_identity_override_document_key( - feature_id=feature_state.feature.id, - identity_uuid=identity_uuid, - ), - environment_id=str(environment_id), - environment_api_key=environment_api_key, - feature_state=feature_state, - identifier=identifier, - identity_uuid=identity_uuid, + return map_identity_override_document_to_identity_override( + { + "environment_id": str(environment_id), + "document_key": get_environments_v2_identity_override_document_key( + feature_id=int(feature_state["feature"]["id"]), + identity_uuid=identity_uuid, + ), + "environment_api_key": environment_api_key, + "identifier": identifier, + "identity_uuid": identity_uuid, + "feature_state": feature_state, + "created_date": timezone.now(), + } ) @@ -167,10 +193,9 @@ def map_identity_changeset_to_identity_override_changeset( for _, change_details in identity_changeset["feature_overrides"].items(): match change_details["change_type"]: case "-": - feature_state = FeatureStateModel.parse_obj(change_details["old"]) to_delete.append( map_engine_feature_state_to_identity_override( - feature_state=feature_state, + feature_state=change_details["old"], identity_uuid=identity_uuid, identifier=identifier, environment_api_key=environment_api_key, @@ -178,10 +203,9 @@ def map_identity_changeset_to_identity_override_changeset( ) ) case _: - feature_state = FeatureStateModel.parse_obj(change_details["new"]) to_put.append( map_engine_feature_state_to_identity_override( - feature_state=feature_state, + feature_state=change_details["new"], identity_uuid=identity_uuid, identifier=identifier, environment_api_key=environment_api_key, @@ -195,10 +219,23 @@ def map_identity_changeset_to_identity_override_changeset( def map_identity_override_to_identity_override_document( identity_override: IdentityOverrideV2, ) -> Document: - return { - field_name: _map_value_to_document_value(value) - for field_name, value in identity_override - } + return cast(Document, identity_override) + + +def _validate_document(adapter: TypeAdapter[T], document: Mapping[str, Any]) -> T: + # The schema doesn't round-trip stored numbers, e.g. it validates an integer + # `Decimal` feature value as a string, so they're validated as native numbers. + return adapter.validate_python(_map_decimals_to_numbers(document)) + + +def _map_decimals_to_numbers(value: Any) -> Any: + if isinstance(value, Decimal): + return float(value) if value.as_tuple().exponent else int(value) + if isinstance(value, Mapping): + return {key: _map_decimals_to_numbers(item) for key, item in value.items()} + if isinstance(value, list): + return [_map_decimals_to_numbers(item) for item in value] + return value def _get_compressed_environment_document( @@ -214,55 +251,3 @@ def _get_compressed_environment_document( compressed_size_bytes=compressed_size_bytes, compression_ratio=compressed_size_bytes / uncompressed_size_bytes, ) - - -T = TypeVar("T") - - -def _noop_encoder(value: T) -> T: - return value - - -def _base_model_encoder(value: BaseModel) -> DocumentValue: - return _map_value_to_document_value(value.dict()) - - -def _dict_encoder(value: Dict[str, Any]) -> Dict[str, DocumentValue]: - return {f_name: _map_value_to_document_value(val) for f_name, val in value.items()} - - -def _list_encoder(value: List[Any]) -> List[DocumentValue]: - return [_map_value_to_document_value(item) for item in value] - - -def _decimal_encoder(value: Union[int, float]) -> Decimal: - return Decimal(str(value)) - - -def _isoformat_encoder(value: datetime) -> str: - return value.isoformat() - - -DOCUMENT_VALUE_ENCODERS_BY_TYPE: dict[type, Callable[[Any], DocumentValue]] = { - BaseModel: _base_model_encoder, - dict: _dict_encoder, - list: _list_encoder, - type(None): _noop_encoder, - str: _noop_encoder, - bool: _noop_encoder, - int: _decimal_encoder, - float: _decimal_encoder, - datetime: _isoformat_encoder, - Decimal: _noop_encoder, -} - - -def _map_value_to_document_value(value: Any) -> DocumentValue: - for base in value.__class__.__mro__[:-1]: - try: - encoder = DOCUMENT_VALUE_ENCODERS_BY_TYPE[base] - except KeyError: - continue - return encoder(value) - else: - return str(value) diff --git a/api/util/mappers/engine.py b/api/util/mappers/engine.py index b3f07d28b7ec..b4f87f5e0332 100644 --- a/api/util/mappers/engine.py +++ b/api/util/mappers/engine.py @@ -1,11 +1,13 @@ +import uuid from collections.abc import Iterable from itertools import chain from typing import TYPE_CHECKING, Any, Dict, List, Optional from uuid import UUID +from django.utils import timezone + from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES from features.versioning.models import EnvironmentFeatureVersion -from util.engine_models.identities.models import IdentityModel if TYPE_CHECKING: # pragma: no cover from environments.identities.models import ( # type: ignore[attr-defined] @@ -29,6 +31,7 @@ "map_environment_api_key_to_engine", "map_environment_to_engine", "map_feature_to_engine", + "map_identifier_to_engine", "map_identity_to_engine", "map_mv_option_to_engine", "map_segment_to_engine", @@ -365,7 +368,7 @@ def map_identity_to_engine( *, with_overrides: bool = True, with_traits: bool = True, -) -> IdentityModel: +) -> dict[str, Any]: environment_api_key = identity.environment.api_key # Read relationships - grab all the data needed from the ORM here. @@ -397,21 +400,37 @@ def map_identity_to_engine( ] identity_trait_models = map_traits_to_engine(identity_traits) - return IdentityModel.model_validate( - { - # Attributes: - "identifier": identity.identifier, - "environment_api_key": environment_api_key, - "created_date": identity.created_date, - "django_id": identity.pk, - # - # Relationships: - "identity_features": identity_feature_state_models, - "identity_traits": identity_trait_models, - } + return map_identifier_to_engine( + identity.identifier, + environment_api_key, + created_date=identity.created_date, + identity_features=identity_feature_state_models, + identity_traits=identity_trait_models, + django_id=identity.pk, ) +def map_identifier_to_engine( + identifier: str, + environment_api_key: str, + **fields: Any, +) -> dict[str, Any]: + """An identity's document fields, defaulted as for a new identity.""" + return { + "identifier": identifier, + "environment_api_key": environment_api_key, + "created_date": timezone.now(), + "identity_features": [], + "identity_traits": [], + "system_traits": None, + "identity_uuid": uuid.uuid4(), + "django_id": None, + "dashboard_alias": None, + "composite_key": f"{environment_api_key}_{identifier}", + **fields, + } + + def _get_prioritised_feature_states( feature_states: Iterable["FeatureState"], ) -> List["FeatureState"]: diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py index 80b82e2aea29..5a32a27c7c0a 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -39,9 +39,13 @@ def map_environment_to_sdk_document(environment: "Environment") -> SDKDocument: identities_with_overrides[identity_id] = feature_state.identity engine_environment["identity_overrides"] = [ # System-owned identity data must never reach local-eval SDKs. - map_identity_to_engine(identity, with_traits=False).model_dump( - exclude={"system_traits"} - ) + { + field_name: value + for field_name, value in map_identity_to_engine( + identity, with_traits=False + ).items() + if field_name != "system_traits" + } for identity in identities_with_overrides.values() ] From 0ac09ae6009300af0b4586ef08fbe0e9458447ff Mon Sep 17 00:00:00 2001 From: Kim Gustyr Date: Sat, 26 Sep 2026 22:54:19 +0200 Subject: [PATCH 4/5] chore(api): take the stored Decimals fix from the flagsmith-common branch Drops the Core-side workaround, now that flagsmith-common keeps stored Decimals when validating documents. --- api/pyproject.toml | 1 + .../mappers/test_unit_mappers_dynamodb.py | 50 +------------------ api/util/mappers/dynamodb.py | 27 ++-------- api/uv.lock | 10 ++-- 4 files changed, 9 insertions(+), 79 deletions(-) diff --git a/api/pyproject.toml b/api/pyproject.toml index f740df3dec32..86bf18d9772a 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -153,6 +153,7 @@ explicit = true [tool.uv.sources] flagsmith-private = { index = "flagsmith-pypi-production" } +flagsmith-common = { git = "https://github.com/Flagsmith/flagsmith-common", branch = "fix/schemas-stored-decimal-round-trip" } [tool.uv] required-version = ">=0.11.18" # Ensure this matches the version in .pre-commit-config.yaml diff --git a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py index b6a328026a3d..c5d8dca15026 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py @@ -2,7 +2,7 @@ import json import uuid from decimal import Decimal -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from django.utils import timezone @@ -201,54 +201,6 @@ def test_identity_document__system_traits_set__round_trip_preserves_system_trait assert parsed["system_traits"] == {"flagsmith_cohort_2b6d1f5f": True} -def test_map_engine_identity_to_identity_document__stored_numbers__round_trip_unchanged() -> ( - None -): - # Given - stored_document = dynamodb.map_engine_identity_to_identity_document( - map_identifier_to_engine( - "test_identity", - "api-key", - identity_traits=[ - {"trait_key": "integer", "trait_value": 1}, - {"trait_key": "float", "trait_value": 1.5}, - ], - identity_features=[ - { - "feature": {"id": 1, "name": "feature", "type": "MULTIVARIATE"}, - "enabled": True, - "feature_state_value": 5, - "featurestate_uuid": str(uuid.uuid4()), - "multivariate_feature_state_values": [ - { - "mv_fs_value_uuid": str(uuid.uuid4()), - "percentage_allocation": 100, - "multivariate_feature_option": {"id": 2, "value": 3}, - } - ], - } - ], - ) - ) - - # When - document: dict[str, Any] = dynamodb.map_engine_identity_to_identity_document( - dynamodb.map_identity_document_to_engine_identity(stored_document) - ) - - # Then - assert document == stored_document - assert document["identity_traits"] == [ - {"trait_key": "integer", "trait_value": Decimal("1")}, - {"trait_key": "float", "trait_value": Decimal("1.5")}, - ] - (feature_state,) = document["identity_features"] - assert feature_state["feature_state_value"] == Decimal("5") - assert feature_state["multivariate_feature_state_values"][0][ - "multivariate_feature_option" - ]["value"] == Decimal("3") - - def test_map_environment_to_environment_v2_document__valid_environment__returns_expected_document( environment: "Environment", feature_state: "FeatureState", diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index 61adecc054fd..adbfd618bae7 100644 --- a/api/util/mappers/dynamodb.py +++ b/api/util/mappers/dynamodb.py @@ -1,6 +1,5 @@ from collections.abc import Mapping -from decimal import Decimal -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, cast from django.utils import timezone from flagsmith_schemas.dynamodb import ( @@ -53,8 +52,6 @@ ) -T = TypeVar("T") - _environment_adapter: TypeAdapter[Environment] = TypeAdapter(Environment) _environment_api_key_adapter: TypeAdapter[EnvironmentAPIKey] = TypeAdapter( EnvironmentAPIKey @@ -127,14 +124,14 @@ def map_environment_api_key_to_environment_api_key_document( def map_identity_document_to_engine_identity( identity_document: Mapping[str, Any], ) -> Identity: - return _validate_document(_identity_adapter, identity_document) + return _identity_adapter.validate_python(identity_document) def map_engine_identity_to_identity_document( engine_identity: Mapping[str, Any], ) -> Document: identity_document = cast( - Document, _validate_document(_identity_adapter, engine_identity) + Document, _identity_adapter.validate_python(engine_identity) ) return { field_name: value @@ -152,7 +149,7 @@ def map_identity_to_identity_document( def map_identity_override_document_to_identity_override( identity_override_document: Mapping[str, Any], ) -> IdentityOverrideV2: - return _validate_document(_identity_override_adapter, identity_override_document) + return _identity_override_adapter.validate_python(identity_override_document) def map_engine_feature_state_to_identity_override( @@ -222,22 +219,6 @@ def map_identity_override_to_identity_override_document( return cast(Document, identity_override) -def _validate_document(adapter: TypeAdapter[T], document: Mapping[str, Any]) -> T: - # The schema doesn't round-trip stored numbers, e.g. it validates an integer - # `Decimal` feature value as a string, so they're validated as native numbers. - return adapter.validate_python(_map_decimals_to_numbers(document)) - - -def _map_decimals_to_numbers(value: Any) -> Any: - if isinstance(value, Decimal): - return float(value) if value.as_tuple().exponent else int(value) - if isinstance(value, Mapping): - return {key: _map_decimals_to_numbers(item) for key, item in value.items()} - if isinstance(value, list): - return [_map_decimals_to_numbers(item) for item in value] - return value - - def _get_compressed_environment_document( document: Document, adapter: "TypeAdapter[Any]", diff --git a/api/uv.lock b/api/uv.lock index 248d7f8439e5..6957cc6360b3 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1470,8 +1470,8 @@ requires-dist = [ { name = "drf-writable-nested", specifier = ">=0.6.2,<0.7.0" }, { name = "environs", specifier = ">=14.1.1,<15.0.0" }, { name = "flagsmith", specifier = ">=6.2.0,<7.0.0" }, - { name = "flagsmith-common", extras = ["common-core", "flagsmith-schemas", "task-processor"], specifier = ">=3.15.0,<4" }, - { name = "flagsmith-common", extras = ["test-tools"], marker = "extra == 'dev'" }, + { name = "flagsmith-common", extras = ["common-core", "flagsmith-schemas", "task-processor"], git = "https://github.com/Flagsmith/flagsmith-common?branch=fix%2Fschemas-stored-decimal-round-trip" }, + { name = "flagsmith-common", extras = ["test-tools"], marker = "extra == 'dev'", git = "https://github.com/Flagsmith/flagsmith-common?branch=fix%2Fschemas-stored-decimal-round-trip" }, { name = "flagsmith-flag-engine", specifier = ">=11.1.0,<12.0.0" }, { name = "flagsmith-private", marker = "extra == 'private'", specifier = ">=0.14.0,<1", index = "https://flagsmith-production-084060095745.d.codeartifact.eu-west-2.amazonaws.com/pypi/flagsmith-pypi-production/simple/" }, { name = "flagsmith-sql-flag-engine", specifier = ">=0.2.0,<0.3.0" }, @@ -1543,11 +1543,7 @@ provides-extras = ["private", "dev"] [[package]] name = "flagsmith-common" version = "3.15.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/59/a4/c8605bb5fbf853a8d06838241d89d0b70ea93b5c5b68c234f51c660b74b0/flagsmith_common-3.15.0.tar.gz", hash = "sha256:970290aa8b5aa0399c6e9f64b843de95e5a458fcbe0623425a6d15305c7e3559", size = 62639, upload-time = "2026-09-15T11:15:51.097Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/97/57/08475196dc972e2d980b267d0bea25353c402983a8ea13f9fdfed6745f38/flagsmith_common-3.15.0-py3-none-any.whl", hash = "sha256:bd7b54970fb0ebe563fe007ef25cbc45bd1a4c586648d8c2bf6df27e82469565", size = 100743, upload-time = "2026-09-15T11:15:49.505Z" }, -] +source = { git = "https://github.com/Flagsmith/flagsmith-common?branch=fix%2Fschemas-stored-decimal-round-trip#8aa120be0c70e70e260dc9fddd6dd9df267d76dc" } [package.optional-dependencies] common-core = [ From 84f9bcdd8f73f6f307202e3fc61245829463121b Mon Sep 17 00:00:00 2001 From: "flagsmith-engineering[bot]" Date: Sat, 26 Sep 2026 20:55:26 +0000 Subject: [PATCH 5/5] chore: Update documentation artefacts --- .../observability/_events-catalogue.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index 6001fc751eb0..5710f42a303e 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md @@ -623,7 +623,7 @@ Attributes: ### `segment_membership.compute.segment.skipped` Logged at `error` from: - - `api/segment_membership/services.py:149` + - `api/segment_membership/services.py:148` Attributes: - `project.id` @@ -633,7 +633,7 @@ Attributes: ### `segment_membership.members.segment.skipped` Logged at `error` from: - - `api/segment_membership/services.py:215` + - `api/segment_membership/services.py:214` Attributes: - `reason`