diff --git a/flagsmith/flagsmith.py b/flagsmith/flagsmith.py index d8f34b2..ee028a2 100644 --- a/flagsmith/flagsmith.py +++ b/flagsmith/flagsmith.py @@ -15,28 +15,20 @@ EventProcessorConfig, ) from flagsmith.exceptions import FlagsmithAPIError, FlagsmithClientError +from flagsmith.local_evaluation import LocalEvaluation from flagsmith.mappers import ( map_context_and_identity_data_to_context, - map_environment_document_to_context, map_environment_document_to_environment_updated_at, map_segment_results_to_identity_segments, resolve_trait_values, ) -from flagsmith.models import ( - DefaultFlag, - Flag, - Flags, - Segment, - SegmentOverridesIndex, - build_segment_overrides_index, -) +from flagsmith.models import DefaultFlag, Flag, Flags, Segment from flagsmith.offline_handlers import OfflineHandler from flagsmith.polling_manager import EnvironmentDataPollingManager from flagsmith.streaming_manager import EventStreamManager from flagsmith.types import ( ApplicationMetadata, JsonType, - SDKEvaluationContext, StreamEvent, TraitMapping, ) @@ -135,8 +127,7 @@ def __init__( self.enable_realtime_updates = enable_realtime_updates self._analytics_processor: typing.Optional[AnalyticsProcessor] = None self._event_processor: typing.Optional[EventProcessor] = None - self.__evaluation_context: typing.Optional[SDKEvaluationContext] = None - self._segment_overrides_index: SegmentOverridesIndex = {} + self._local_evaluation: typing.Optional[LocalEvaluation] = None self._environment_updated_at: typing.Optional[datetime] = None # argument validation @@ -158,7 +149,7 @@ def __init__( ) if self.offline_handler: - self._evaluation_context = map_environment_document_to_context( + self._local_evaluation = LocalEvaluation.from_environment_document( self.offline_handler.get_environment() ) @@ -248,12 +239,13 @@ def _initialise_local_evaluation(self) -> None: # method calls, update the environment manually. self.update_environment() if self.enable_realtime_updates: - if not self._evaluation_context: + if not self._local_evaluation: raise ValueError("Unable to get environment from API key") + environment_key = self._local_evaluation.context["environment"]["key"] stream_url = urljoin( self.realtime_api_url, - f"sse/environments/{self._evaluation_context['environment']['key']}/stream", + f"sse/environments/{environment_key}/stream", ) self.event_stream_thread = EventStreamManager( @@ -291,7 +283,7 @@ def get_environment_flags(self) -> Flags: """ if ( self.offline_mode or self.enable_local_evaluation - ) and self._evaluation_context: + ) and self._local_evaluation: return self._get_environment_flags_from_document() return self._get_environment_flags_from_api() @@ -319,7 +311,7 @@ def get_identity_flags( traits = traits or {} if ( self.offline_mode or self.enable_local_evaluation - ) and self._evaluation_context: + ) and self._local_evaluation: return self._get_identity_flags_from_document(identifier, traits) return self._get_identity_flags_from_api( identifier, @@ -341,13 +333,13 @@ def get_identity_segments( Flagsmith, e.g. {"num_orders": 10} :return: list of Segment objects that the identity is part of. """ - if not self._evaluation_context: + if not (local_evaluation := self._local_evaluation): raise FlagsmithClientError( "Local evaluation required to obtain identity segments." ) context = map_context_and_identity_data_to_context( - context=self._evaluation_context, + context=local_evaluation.context, identifier=identifier, traits=traits, ) @@ -470,8 +462,8 @@ def update_environment(self) -> None: logger.exception("Error retrieving environment document from API") else: try: - self._evaluation_context = map_environment_document_to_context( - environment_data, + self._local_evaluation = LocalEvaluation.from_environment_document( + environment_data ) self._environment_updated_at = ( map_environment_document_to_environment_updated_at( @@ -481,26 +473,6 @@ def update_environment(self) -> None: except (KeyError, TypeError, ValueError): logger.exception("Error parsing environment document") - @property - def _evaluation_context(self) -> typing.Optional[SDKEvaluationContext]: - return self.__evaluation_context - - @_evaluation_context.setter - def _evaluation_context( - self, context: typing.Optional[SDKEvaluationContext] - ) -> None: - """Swap in a new evaluation context and rebuild the overrides index. - - The index maps feature_name -> segments that override it. Built once - per refresh and reused across every subsequent per-identity lazy - resolution; rebuilding here keeps it in sync with the current doc - without any hot-path cost. - """ - self.__evaluation_context = context - self._segment_overrides_index = ( - build_segment_overrides_index(context) if context is not None else {} - ) - def _get_headers( self, environment_key: str, @@ -520,22 +492,13 @@ def _get_headers( return headers def _get_environment_flags_from_document(self) -> Flags: - if self._evaluation_context is None: + if (local_evaluation := self._local_evaluation) is None: raise TypeError("No environment present") - # Omit segments from evaluation context for environment flags - # as they are only relevant for identity-specific evaluations - context_without_segments = self._evaluation_context.copy() - context_without_segments.pop("segments", None) - - evaluation_result = engine.get_evaluation_result( - context=context_without_segments, - ) - - return Flags.from_evaluation_result( - evaluation_result=evaluation_result, - analytics_processor=self._analytics_processor, + return Flags( + flags=dict(local_evaluation.environment_flags), default_flag_handler=self.default_flag_handler, + _analytics_processor=self._analytics_processor, ) def _get_identity_flags_from_document( @@ -543,11 +506,10 @@ def _get_identity_flags_from_document( identifier: str, traits: TraitMapping, ) -> Flags: - if self._evaluation_context is None: + if (local_evaluation := self._local_evaluation) is None: raise TypeError("No environment present") - context = map_context_and_identity_data_to_context( - context=self._evaluation_context, + context, overrides_index = local_evaluation.get_identity_context( identifier=identifier, traits=traits, ) @@ -556,7 +518,7 @@ def _get_identity_flags_from_document( # large environment. return Flags.from_evaluation_context( context=context, - overrides_index=self._segment_overrides_index, + overrides_index=overrides_index, analytics_processor=self._analytics_processor, default_flag_handler=self.default_flag_handler, ) diff --git a/flagsmith/local_evaluation.py b/flagsmith/local_evaluation.py new file mode 100644 index 0000000..1bfc632 --- /dev/null +++ b/flagsmith/local_evaluation.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import typing +from collections import ChainMap +from functools import cached_property + +from flag_engine import engine +from flag_engine.context.types import SegmentContext + +from flagsmith.api.types import EnvironmentModel +from flagsmith.mappers import ( + map_context_and_identity_data_to_context, + map_environment_document_to_context, + map_environment_document_to_identity_overrides_segments, +) +from flagsmith.models import ( + Flag, + SegmentOverridesIndex, + build_segment_overrides_index, +) +from flagsmith.types import ( + FeatureMetadata, + SDKEvaluationContext, + SegmentMetadata, + TraitMapping, +) + + +class LocalEvaluation: + """Per-environment-document state for local evaluation.""" + + def __init__( + self, + context: SDKEvaluationContext, + identity_overrides_segments: typing.Optional[ + typing.Mapping[str, SegmentContext[SegmentMetadata, FeatureMetadata]] + ] = None, + ) -> None: + self.context = context + self.overrides_index: SegmentOverridesIndex = build_segment_overrides_index( + context + ) + self._identity_overrides = { + identifier: (segment_context, self._index_overrides(segment_context)) + for identifier, segment_context in ( + identity_overrides_segments or {} + ).items() + } + + @classmethod + def from_environment_document( + cls, + environment_document: EnvironmentModel, + ) -> LocalEvaluation: + return cls( + context=map_environment_document_to_context(environment_document), + identity_overrides_segments=( + map_environment_document_to_identity_overrides_segments( + environment_document + ) + ), + ) + + @cached_property + def environment_flags(self) -> typing.Dict[str, Flag]: + """Environment flags, which depend on the environment document only.""" + context_without_segments = self.context.copy() + context_without_segments.pop("segments", None) + evaluation_result = engine.get_evaluation_result( + context=context_without_segments, + ) + return { + feature_name: Flag.from_evaluation_result(flag_result) + for feature_name, flag_result in evaluation_result["flags"].items() + } + + def get_identity_context( + self, + identifier: str, + traits: typing.Optional[TraitMapping] = None, + ) -> typing.Tuple[SDKEvaluationContext, SegmentOverridesIndex]: + """Evaluation context and overrides index for an identity, including + its own identity overrides segment, if any.""" + context = map_context_and_identity_data_to_context( + context=self.context, + identifier=identifier, + traits=traits, + ) + if (identity_overrides := self._identity_overrides.get(identifier)) is None: + return context, self.overrides_index + + segment_context, identity_overrides_index = identity_overrides + context["segments"] = { + **(context.get("segments") or {}), + segment_context["key"]: segment_context, + } + return context, ChainMap(identity_overrides_index, self.overrides_index) + + def _index_overrides( + self, + segment_context: SegmentContext[SegmentMetadata, FeatureMetadata], + ) -> SegmentOverridesIndex: + """The overrides index, for the features `segment_context` overrides.""" + return { + feature_name: { + **(self.overrides_index.get(feature_name) or {}), + segment_context["key"]: segment_context, + } + for override in segment_context.get("overrides") or () + if (feature_name := override["name"]) + } diff --git a/flagsmith/mappers.py b/flagsmith/mappers.py index 091cf3e..7e22a9b 100644 --- a/flagsmith/mappers.py +++ b/flagsmith/mappers.py @@ -1,13 +1,13 @@ import json import typing import uuid -from collections import defaultdict from datetime import datetime, timezone import sseclient from flag_engine.context.types import ( FeatureContext, FeatureValue, + InOperatorSegmentCondition, SegmentContext, SegmentRule, StrValueSegmentCondition, @@ -17,7 +17,6 @@ from flagsmith.api.types import ( EnvironmentModel, FeatureStateModel, - IdentityModel, SegmentRuleModel, ) from flagsmith.models import Segment @@ -30,13 +29,7 @@ ) from flagsmith.utils.datetime import fromisoformat -OverrideKey = typing.Tuple[ - int, - str, - bool, - typing.Any, -] -OverridesKey = typing.Tuple[OverrideKey, ...] +IDENTITY_OVERRIDES_SEGMENT_KEY = "identity_overrides" def map_segment_results_to_identity_segments( @@ -115,96 +108,68 @@ def map_environment_document_to_context( ) }, "segments": { - **{ - (segment_key := str(segment_id := segment["id"])): { - "key": segment_key, - "name": segment["name"], - "rules": _map_environment_document_rules_to_context_rules( - segment["rules"] - ), - "overrides": list( - _map_environment_document_feature_states_to_feature_contexts( - segment.get("feature_states") or [] - ) - ), - "metadata": SegmentMetadata( - id=segment_id, - source="api", - ), - } - for segment in environment_document["project"]["segments"] - }, - **_map_identity_overrides_to_segments( - environment_document.get("identity_overrides") or [] - ), + (segment_key := str(segment_id := segment["id"])): { + "key": segment_key, + "name": segment["name"], + "rules": _map_environment_document_rules_to_context_rules( + segment["rules"] + ), + "overrides": list( + _map_environment_document_feature_states_to_feature_contexts( + segment.get("feature_states") or [] + ) + ), + "metadata": SegmentMetadata( + id=segment_id, + source="api", + ), + } + for segment in environment_document["project"]["segments"] }, } -def _map_identity_overrides_to_segments( - identity_overrides: list[IdentityModel], -) -> dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]: - features_to_identifiers: typing.Dict[ - OverridesKey, - typing.List[str], - ] = defaultdict(list) - for identity_override in identity_overrides: - identity_features = identity_override["identity_features"] - if not identity_features: - continue - overrides_key = tuple( - ( - feature_state["feature"]["id"], - feature_state["feature"]["name"], - feature_state["enabled"], - feature_state["feature_state_value"], - ) - for feature_state in sorted( - identity_features, - key=lambda feature_state: feature_state["feature"]["name"], - ) - ) - features_to_identifiers[overrides_key].append(identity_override["identifier"]) - segment_contexts: typing.Dict[ - str, - SegmentContext[ - SegmentMetadata, - FeatureMetadata, - ], - ] = {} - for overrides_key, identifiers in features_to_identifiers.items(): - # Create a segment context for each unique set of overrides - # Generate a unique key to avoid collisions - segment_key = str(hash(overrides_key)) - segment_contexts[segment_key] = SegmentContext( - key=segment_key, - name="identity_overrides", - rules=[ - { - "type": "ALL", - "conditions": [ - { - "property": "$.identity.identifier", - "operator": "IN", - "value": identifiers, - } - ], - } +def map_environment_document_to_identity_overrides_segments( + environment_document: EnvironmentModel, +) -> typing.Dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]]: + """Map identity overrides to segment contexts by identifier, each matching + its identity alone and carrying its overrides.""" + return { + identity_override["identifier"]: { + "key": IDENTITY_OVERRIDES_SEGMENT_KEY, + "name": "identity_overrides", + "rules": [ + map_identifiers_to_identity_overrides_rule( + [identity_override["identifier"]] + ) ], - overrides=[ + "overrides": [ { "key": "", # Identity overrides never carry multivariate options - "name": feature_name, - "enabled": feature_enabled, - "value": feature_value, + "name": feature_state["feature"]["name"], + "enabled": feature_state["enabled"], + "value": feature_state["feature_state_value"], "priority": float("-inf"), # Highest possible priority - "metadata": {"id": feature_id}, + "metadata": {"id": feature_state["feature"]["id"]}, } - for feature_id, feature_name, feature_enabled, feature_value in overrides_key + for feature_state in identity_features ], - metadata=SegmentMetadata(source="identity_overrides"), - ) - return segment_contexts + "metadata": SegmentMetadata(source="identity_overrides"), + } + for identity_override in environment_document.get("identity_overrides") or [] + if (identity_features := identity_override["identity_features"]) + } + + +def map_identifiers_to_identity_overrides_rule( + identifiers: typing.List[str], +) -> SegmentRule: + condition: InOperatorSegmentCondition = { + "property": "$.identity.identifier", + "operator": "IN", + "value": identifiers, + } + return {"type": "ALL", "conditions": [condition]} def _map_environment_document_rules_to_context_rules( diff --git a/flagsmith/models.py b/flagsmith/models.py index d29f86c..d4880cf 100644 --- a/flagsmith/models.py +++ b/flagsmith/models.py @@ -16,9 +16,12 @@ SegmentMetadata, ) -SegmentOverridesIndex = typing.Dict[ - str, typing.Dict[str, SegmentContext[SegmentMetadata, FeatureMetadata]] +SegmentKey = str +FeatureName = str +SegmentContexts = typing.Dict[ + SegmentKey, SegmentContext[SegmentMetadata, FeatureMetadata] ] +SegmentOverridesIndex = typing.MutableMapping[FeatureName, SegmentContexts] def build_segment_overrides_index( diff --git a/tests/test_flagsmith.py b/tests/test_flagsmith.py index 9296feb..b3cbfdd 100644 --- a/tests/test_flagsmith.py +++ b/tests/test_flagsmith.py @@ -17,6 +17,7 @@ FlagsmithAPIError, FlagsmithFeatureDoesNotExistError, ) +from flagsmith.local_evaluation import LocalEvaluation from flagsmith.models import DefaultFlag, ExperimentMetadata, Flag, Flags from flagsmith.offline_handlers import OfflineHandler from flagsmith.types import SDKEvaluationContext @@ -47,14 +48,16 @@ def test_update_environment_sets_environment( ) -> None: # Given responses.add(method="GET", url=flagsmith.environment_url, body=environment_json) - assert flagsmith._evaluation_context is None + assert flagsmith._local_evaluation is None # When flagsmith.update_environment() # Then - assert flagsmith._evaluation_context is not None - assert flagsmith._evaluation_context == evaluation_context + assert flagsmith._local_evaluation is not None + assert flagsmith._local_evaluation.context == ( + LocalEvaluation(evaluation_context).context + ) @responses.activate() @@ -84,7 +87,7 @@ def test_get_environment_flags_uses_local_environment_when_available( evaluation_context: SDKEvaluationContext, ) -> None: # Given - flagsmith._evaluation_context = evaluation_context + flagsmith._local_evaluation = LocalEvaluation(evaluation_context) flagsmith.enable_local_evaluation = True # When @@ -134,6 +137,41 @@ def test_get_environment_flags_omits_segments_from_evaluation_context( mock_get_evaluation_result.assert_called_once_with(context=context_without_segments) +def test_get_environment_flags__local_evaluation__evaluates_once_per_environment( + mocker: MockerFixture, + local_eval_flagsmith: Flagsmith, +) -> None: + # Given + spy = mocker.spy(engine, "get_evaluation_result") + + # When + first = local_eval_flagsmith.get_environment_flags() + second = local_eval_flagsmith.get_environment_flags() + + # Then + # the environment is evaluated once, each call getting its own mapping + assert spy.call_count == 1 + assert first.flags == second.flags + assert first.flags is not second.flags + + +def test_get_environment_flags__environment_refreshed__evaluates_again( + mocker: MockerFixture, + local_eval_flagsmith: Flagsmith, + evaluation_context: SDKEvaluationContext, +) -> None: + # Given + local_eval_flagsmith.get_environment_flags() + spy = mocker.spy(engine, "get_evaluation_result") + + # When + local_eval_flagsmith._local_evaluation = LocalEvaluation(evaluation_context) + local_eval_flagsmith.get_environment_flags() + + # Then + spy.assert_called_once() + + @responses.activate() def test_get_identity_flags_calls_api_when_no_local_environment_no_traits( flagsmith: Flagsmith, identities_json: str @@ -196,7 +234,7 @@ def test_get_identity_flags_uses_local_environment_when_available( mocker: MockerFixture, ) -> None: # Given - flagsmith._evaluation_context = evaluation_context + flagsmith._local_evaluation = LocalEvaluation(evaluation_context) flagsmith.enable_local_evaluation = True # `Flags` materialises identity flags via `engine.get_evaluation_result` # imported from `flagsmith.models`, so patch it where it's actually used. diff --git a/tests/test_local_evaluation.py b/tests/test_local_evaluation.py new file mode 100644 index 0000000..96ae1ea --- /dev/null +++ b/tests/test_local_evaluation.py @@ -0,0 +1,143 @@ +import pytest +from flag_engine import engine +from pytest_mock import MockerFixture + +from flagsmith.api.types import EnvironmentModel +from flagsmith.local_evaluation import LocalEvaluation + + +@pytest.fixture() +def local_evaluation(environment: EnvironmentModel) -> LocalEvaluation: + return LocalEvaluation.from_environment_document(environment) + + +def test_local_evaluation__context__excludes_identity_overrides_segments( + local_evaluation: LocalEvaluation, +) -> None: + # When + segments = local_evaluation.context["segments"] + + # Then + assert segments is not None + assert [segment["name"] for segment in segments.values()] == ["Test segment"] + assert local_evaluation.overrides_index == {} + + +def test_local_evaluation__get_identity_context_no_identity_overrides__returns_context( + local_evaluation: LocalEvaluation, +) -> None: + # When + context, overrides_index = local_evaluation.get_identity_context( + "not-overridden-id", {"foo": "bar"} + ) + + # Then + assert context["segments"] is local_evaluation.context["segments"] + assert overrides_index is local_evaluation.overrides_index + assert context["identity"] == { + "identifier": "not-overridden-id", + "traits": {"foo": "bar"}, + } + + +@pytest.mark.parametrize( + "identifier, expected_value", + [ + ("overridden-id", "some-overridden-value"), + ("another-overridden-id", "another-overridden-value"), + ], +) +def test_local_evaluation__get_identity_context_identity_overrides__returns_own_narrowed( + local_evaluation: LocalEvaluation, + identifier: str, + expected_value: str, +) -> None: + # When + context, overrides_index = local_evaluation.get_identity_context(identifier) + + # Then + # only the identity's own identity overrides segment is added, + # matching its identifier alone + assert context["segments"] is not None + identity_segments = [ + segment + for segment in context["segments"].values() + if segment["name"] == "identity_overrides" + ] + assert identity_segments == [ + { + "key": "identity_overrides", + "name": "identity_overrides", + "rules": [ + { + "type": "ALL", + "conditions": [ + { + "property": "$.identity.identifier", + "operator": "IN", + "value": [identifier], + } + ], + } + ], + "overrides": [ + { + "key": "", + "name": "some_feature", + "enabled": False, + "value": expected_value, + "priority": float("-inf"), + "metadata": {"id": 1}, + } + ], + "metadata": {"source": "identity_overrides"}, + } + ] + [segment] = identity_segments + assert list(overrides_index) == ["some_feature"] + assert list(overrides_index["some_feature"].values()) == [segment] + + # And the shared state is left untouched + assert local_evaluation.overrides_index == {} + assert local_evaluation.context["segments"] is not None + assert len(local_evaluation.context["segments"]) == 1 + + +@pytest.mark.parametrize( + "identifier, expected_value", + [ + ("overridden-id", "some-overridden-value"), + ("another-overridden-id", "another-overridden-value"), + ("not-overridden-id", "some-value"), + ], +) +def test_local_evaluation__get_identity_context__evaluates_to_identity_overrides( + local_evaluation: LocalEvaluation, + identifier: str, + expected_value: str, +) -> None: + # Given + context, _ = local_evaluation.get_identity_context(identifier) + + # When + result = engine.get_evaluation_result(context) + + # Then + assert result["flags"]["some_feature"]["value"] == expected_value + + +def test_local_evaluation__environment_flags__evaluated_once( + local_evaluation: LocalEvaluation, + mocker: MockerFixture, +) -> None: + # Given + spy = mocker.spy(engine, "get_evaluation_result") + + # When + first = local_evaluation.environment_flags + second = local_evaluation.environment_flags + + # Then + assert first is second + assert first["some_feature"].value == "some-value" + spy.assert_called_once() diff --git a/tests/test_mappers.py b/tests/test_mappers.py index 651c010..fd73dca 100644 --- a/tests/test_mappers.py +++ b/tests/test_mappers.py @@ -55,23 +55,3 @@ def test_map_environment_document_to_context__null_variant_key__drops_key() -> N # Then - the null key is dropped, treated as no key variants = context["features"]["mv_feature"]["variants"] assert "key" not in variants[0] - - -def test_map_environment_document_to_context__identity_overrides__keys_are_unique( - environment: EnvironmentModel, -) -> None: - # Given / When - context = map_environment_document_to_context(environment) - - # Then - segments = context["segments"] or {} - identity_override_keys = { - segment_key - for segment_key, segment_context in segments.items() - if segment_context["name"] == "identity_overrides" - } - assert len(identity_override_keys) == 2 - assert all( - segments[segment_key]["key"] == segment_key - for segment_key in identity_override_keys - ) diff --git a/tests/test_streaming_manager.py b/tests/test_streaming_manager.py index 9b2c249..7c43f84 100644 --- a/tests/test_streaming_manager.py +++ b/tests/test_streaming_manager.py @@ -7,6 +7,7 @@ from pytest_mock import MockerFixture from flagsmith import Flagsmith +from flagsmith.local_evaluation import LocalEvaluation from flagsmith.streaming_manager import EventStreamManager from flagsmith.types import StreamEvent @@ -44,7 +45,7 @@ def test_environment_updates_on_recent_event( mocker.patch("flagsmith.Flagsmith.update_environment") flagsmith = Flagsmith(environment_key=server_api_key) - flagsmith._evaluation_context = MagicMock() + flagsmith._local_evaluation = MagicMock(spec=LocalEvaluation) flagsmith._environment_updated_at = environment_updated_at flagsmith.handle_stream_event(event=StreamEvent(updated_at=stream_updated_at)) assert isinstance(flagsmith.update_environment, Mock) @@ -60,7 +61,7 @@ def test_environment_does_not_update_on_past_event( mocker.patch("flagsmith.Flagsmith.update_environment") flagsmith = Flagsmith(environment_key=server_api_key) - flagsmith._evaluation_context = MagicMock() + flagsmith._local_evaluation = MagicMock(spec=LocalEvaluation) flagsmith._environment_updated_at = environment_updated_at flagsmith.handle_stream_event(event=StreamEvent(updated_at=stream_updated_at)) @@ -77,7 +78,7 @@ def test_environment_does_not_update_on_same_event( mocker.patch("flagsmith.Flagsmith.update_environment") flagsmith = Flagsmith(environment_key=server_api_key) - flagsmith._evaluation_context = MagicMock() + flagsmith._local_evaluation = MagicMock(spec=LocalEvaluation) flagsmith._environment_updated_at = environment_updated_at flagsmith.handle_stream_event(event=StreamEvent(updated_at=stream_updated_at))