feat(identities): evaluate system traits when matching segments for edge identities (#8266)

Author: gagantrivedi

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -19,6 +19,9 @@
     CapacityBudgetExceeded,
     SystemTraitWriteRaceError,
 )
+from environments.identities.traits.constants import (
+    TRAIT_STRING_VALUE_MAX_LENGTH,
+)
 from util.engine_models.context.mappers import (
     is_context_in_segment,
     map_environment_identity_to_context,
@@ -103,6 +106,17 @@ def set_system_trait(
         Assumes stored documents never carry `system_traits` as NULL — the
         document mapper omits the attribute when unset.
         """
+        if (
+            isinstance(trait_value, str)
+            and len(trait_value) > TRAIT_STRING_VALUE_MAX_LENGTH
+        ):
+            # The Edge API caps trait values at this length when parsing
+            # identity documents; a longer value there makes every flag
+            # request for the identity fail.
+            raise ValueError(
+                "System trait value must be at most "
+                f"{TRAIT_STRING_VALUE_MAX_LENGTH} characters."
+            )
         composite_key = IdentityModel.generate_composite_key(
             environment_api_key, identifier
         )

api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py

@@ -6,7 +6,7 @@
 from boto3.dynamodb.types import Binary
 from botocore.exceptions import ClientError
 from django.core.exceptions import ObjectDoesNotExist
-from flag_engine.segments.constants import IN
+from flag_engine.segments.constants import IN, IS_SET
 from mypy_boto3_dynamodb.service_resource import Table
 from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
@@ -24,6 +24,9 @@
     SystemTraitWriteRaceError,
 )
 from environments.identities.models import Identity
+from environments.identities.traits.constants import (
+    TRAIT_STRING_VALUE_MAX_LENGTH,
+)
 from environments.identities.traits.models import Trait
 from features.models import Feature, FeatureSegment, FeatureState
 from features.multivariate.models import (
@@ -400,6 +403,55 @@ def test_get_segment_ids__segment_with_feature_overrides__returns_correct_ids(
     assert segment_ids == [identity_matching_segment.id]
 
 
+def test_get_segment_ids__system_trait_backed_segment__returns_correct_ids(
+    project: "Project",
+    environment: "Environment",
+    identity: "Identity",
+    mocker: "MockerFixture",
+) -> None:
+    # Given - two IS_SET segments: one keyed to a system trait the identity
+    # carries, one keyed to a system trait it does not
+    member_segment = Segment.objects.create(name="Cohort members", project=project)
+    rule = SegmentRule.objects.create(segment=member_segment, type=SegmentRule.ALL_RULE)
+    Condition.objects.create(rule=rule, operator=IS_SET, property="flagsmith_cohort_a")
+    other_segment = Segment.objects.create(name="Other cohort", project=project)
+    other_rule = SegmentRule.objects.create(
+        segment=other_segment, type=SegmentRule.ALL_RULE
+    )
+    Condition.objects.create(
+        rule=other_rule, operator=IS_SET, property="flagsmith_cohort_b"
+    )
+
+    identity_document = map_identity_to_identity_document(identity)
+    identity_document["system_traits"] = {"flagsmith_cohort_a": True}
+    identity_uuid = identity_document["identity_uuid"]
+
+    dynamo_identity_wrapper = DynamoIdentityWrapper()
+    mocker.patch.object(
+        dynamo_identity_wrapper, "get_item_from_uuid", return_value=identity_document
+    )
+
+    # When
+    segment_ids = dynamo_identity_wrapper.get_segment_ids(identity_uuid)  # type: ignore[arg-type]
+
+    # Then
+    assert segment_ids == [member_segment.id]
+
+
+def test_set_system_trait__oversized_string_value__raises() -> None:
+    # Given
+    wrapper = DynamoIdentityWrapper()
+
+    # When / Then
+    with pytest.raises(ValueError):
+        wrapper.set_system_trait(
+            environment_api_key="key",
+            identifier="user",
+            trait_key="flagsmith_cohort_a",
+            trait_value="x" * (TRAIT_STRING_VALUE_MAX_LENGTH + 1),
+        )
+
+
 def test_get_segment_ids__in_operator_with_integer_traits__returns_matching_segment(
     project: "Project", environment: "Environment", mocker: "MockerFixture"
 ) -> None:

api/tests/unit/util/engine_models/context/test_unit_context_mappers.py

@@ -0,0 +1,32 @@
+from environments.models import Environment
+from util.engine_models.context.mappers import map_environment_identity_to_context
+from util.engine_models.identities.models import IdentityModel
+from util.engine_models.identities.traits.models import TraitModel
+
+
+def test_map_environment_identity_to_context__system_traits__merged_with_system_winning(
+    environment: Environment,
+) -> None:
+    # Given
+    identity = IdentityModel(
+        identifier="user-1",
+        environment_api_key=environment.api_key,
+        identity_traits=[
+            TraitModel(trait_key="plan", trait_value="free"),
+            TraitModel(trait_key="flagsmith_cohort_a", trait_value="user-written"),
+        ],
+        system_traits={"flagsmith_cohort_a": True},
+    )
+
+    # When
+    context = map_environment_identity_to_context(
+        environment=environment, identity=identity, override_traits=None
+    )
+
+    # Then
+    identity_context = context["identity"]
+    assert identity_context is not None
+    assert identity_context["traits"] == {
+        "plan": "free",
+        "flagsmith_cohort_a": True,
+    }

api/util/engine_models/context/mappers.py

@@ -43,6 +43,16 @@ def map_environment_identity_to_context(
         `identity.identity_traits` if provided.
     :return: An EvaluationContext containing the environment and identity.
     """
+    traits = {
+        trait.trait_key: trait.trait_value
+        for trait in (
+            override_traits if override_traits is not None else identity.identity_traits
+        )
+    }
+    if identity.system_traits:
+        # System-owned traits are not user data: on a key clash, the system
+        # value wins.
+        traits.update(identity.system_traits)
     return {
         "environment": {
             "key": environment.api_key,
@@ -51,14 +61,7 @@ def map_environment_identity_to_context(
         "identity": {
             "identifier": identity.identifier,
             "key": str(identity.django_id or identity.composite_key),
-            "traits": {
-                trait.trait_key: trait.trait_value
-                for trait in (
-                    override_traits
-                    if override_traits is not None
-                    else identity.identity_traits
-                )
-            },
+            "traits": traits,
         },
     }