feat(identities): add system_traits to engine identity document model

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",

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"})
+    ]

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
 

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(

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"}},
 }