From 60bfba4d295cd6ab53ad542700d2964b44b51c5f Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Tue, 4 Aug 2026 10:53:21 +0530 Subject: [PATCH 1/4] feat(identities): add system_traits to engine identity document model --- .../mappers/test_unit_mappers_dynamodb.py | 53 +++++++++++++++++++ api/util/engine_models/identities/models.py | 3 ++ api/util/mappers/dynamodb.py | 2 +- 3 files changed, 57 insertions(+), 1 deletion(-) 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/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( From 4a55ceeb51f4e6905a15a88fb89723bb44964057 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Tue, 4 Aug 2026 11:12:54 +0530 Subject: [PATCH 2/4] fix(identities): exclude system_traits from SDK environment document --- .../util/mappers/test_unit_mappers_sdk.py | 24 +++++++++++++++++++ api/util/mappers/sdk.py | 21 ++++++++++++---- 2 files changed, 40 insertions(+), 5 deletions(-) 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/mappers/sdk.py b/api/util/mappers/sdk.py index 0d9e9865c603..5cc59f381943 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -1,5 +1,7 @@ from typing import TYPE_CHECKING, TypeAlias +from pydantic.main import IncEx + from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES from util.mappers.engine import ( map_environment_to_engine, @@ -10,13 +12,22 @@ 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 | IncEx] = { + **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"}}, } From 3d1fbef1267e89a1b56ba1b9b4ba837976838493 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 05:43:15 +0000 Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- api/util/mappers/sdk.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py index 5cc59f381943..ec757118612d 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -13,12 +13,7 @@ SDKDocumentValue: TypeAlias = ( - dict[str, "SDKDocumentValue"] - | list["SDKDocumentValue"] - | str - | bool - | None - | float + dict[str, "SDKDocumentValue"] | list["SDKDocumentValue"] | str | bool | None | float ) SDKDocument: TypeAlias = dict[str, SDKDocumentValue] From 77dd067775b2955f3324adc0e76f3585f63d5c94 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Wed, 5 Aug 2026 15:49:02 +0530 Subject: [PATCH 4/4] fix(identities): drop private pydantic IncEx import --- api/util/mappers/sdk.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py index 5cc59f381943..26c23ec214f9 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -1,7 +1,5 @@ from typing import TYPE_CHECKING, TypeAlias -from pydantic.main import IncEx - from environments.constants import IDENTITY_INTEGRATIONS_RELATION_NAMES from util.mappers.engine import ( map_environment_to_engine, @@ -13,16 +11,11 @@ SDKDocumentValue: TypeAlias = ( - dict[str, "SDKDocumentValue"] - | list["SDKDocumentValue"] - | str - | bool - | None - | float + dict[str, "SDKDocumentValue"] | list["SDKDocumentValue"] | str | bool | None | float ) SDKDocument: TypeAlias = dict[str, SDKDocumentValue] -SDK_DOCUMENT_EXCLUDE: dict[str, bool | IncEx] = { +SDK_DOCUMENT_EXCLUDE: dict[str, bool | dict[str, set[str]]] = { **dict.fromkeys(IDENTITY_INTEGRATIONS_RELATION_NAMES, True), "dynatrace_config": True, "onboarding_pending": True,