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 32b3f0de2d75..ba2ba78db14c 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.identities.models import IdentityModel from util.mappers import dynamodb from util.mappers.engine import map_feature_state_to_engine @@ -154,6 +155,58 @@ def test_map_identity_to_identity_document__valid_identity__returns_expected_doc assert uuid.UUID(result["identity_uuid"]) # type: ignore[arg-type] +def test_map_engine_identity_to_identity_document__system_traits_set__included_in_document() -> ( + None +): + # Given + engine_identity = IdentityModel( + identifier="test_identity", + environment_api_key="api-key", + system_traits={"flagsmith_cohort_2b6d1f5f": True}, + ) + + # When + result = dynamodb.map_engine_identity_to_identity_document(engine_identity) + + # Then + assert result["system_traits"] == {"flagsmith_cohort_2b6d1f5f": True} + + +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", + ) + + # When + result = dynamodb.map_engine_identity_to_identity_document(engine_identity) + + # Then + assert "system_traits" not in result + + +def test_identity_document__system_traits_set__round_trip_preserves_system_traits() -> ( + None +): + # Given + document = dynamodb.map_engine_identity_to_identity_document( + IdentityModel( + identifier="test_identity", + environment_api_key="api-key", + system_traits={"flagsmith_cohort_2b6d1f5f": True}, + ) + ) + + # When + parsed = IdentityModel.model_validate(document) + + # Then + assert parsed.system_traits == {"flagsmith_cohort_2b6d1f5f": True} + + def test_map_environment_to_environment_v2_document__valid_environment__returns_expected_document( environment: "Environment", feature_state: "FeatureState", 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 825fef2c22d9..aa2af8f3a0d6 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_sdk.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py @@ -3,6 +3,7 @@ import pytest from environments.identities.models import Identity +from util.mappers.engine import map_identity_to_engine from util.mappers.sdk import map_environment_to_sdk_document if TYPE_CHECKING: # pragma: no cover @@ -176,3 +177,26 @@ def test_map_environment_to_sdk_document__identity_overrides_disabled__returns_e "use_identity_composite_key_for_hashing": True, "use_identity_overrides_in_local_eval": False, } + + +def test_map_environment_to_sdk_document__system_traits_set__excluded_from_document( + mocker: "MockerFixture", + environment: "Environment", + identity: Identity, + identity_featurestate: "FeatureState", +) -> None: + # Given + engine_identity = map_identity_to_engine(identity, with_traits=False) + engine_identity.system_traits = {"flagsmith_cohort_2b6d1f5f": True} + mocker.patch( + "util.mappers.sdk.map_identity_to_engine", + return_value=engine_identity, + ) + + # When + result = map_environment_to_sdk_document(environment) + + # Then + assert result["identity_overrides"] == [ + engine_identity.model_dump(exclude={"system_traits"}) + ] diff --git a/api/util/engine_models/identities/models.py b/api/util/engine_models/identities/models.py index ce7173c08309..71518846788e 100644 --- a/api/util/engine_models/identities/models.py +++ b/api/util/engine_models/identities/models.py @@ -7,6 +7,7 @@ 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 @@ -42,6 +43,8 @@ class IdentityModel(BaseModel): 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 diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index dda3ec27bd3e..f59458772629 100644 --- a/api/util/mappers/dynamodb.py +++ b/api/util/mappers/dynamodb.py @@ -53,7 +53,7 @@ TypeAdapter(EnvironmentV2MetaCompressed) ) -_NULLABLE_IDENTITY_KEY_ATTRIBUTES = {"dashboard_alias"} +_NULLABLE_IDENTITY_KEY_ATTRIBUTES = {"dashboard_alias", "system_traits"} def map_environment_to_environment_document( diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py index 0d9e9865c603..26c23ec214f9 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -10,13 +10,17 @@ from environments.models import Environment -SDKDocumentValue: TypeAlias = dict[str, "SDKDocumentValue"] | str | bool | None | float +SDKDocumentValue: TypeAlias = ( + dict[str, "SDKDocumentValue"] | list["SDKDocumentValue"] | str | bool | None | float +) SDKDocument: TypeAlias = dict[str, SDKDocumentValue] -SDK_DOCUMENT_EXCLUDE = { - *IDENTITY_INTEGRATIONS_RELATION_NAMES, - "dynatrace_config", - "onboarding_pending", +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"}}, }