feat(PoC): Segment membership inspection

api/app/settings/common.py

@@ -118,6 +118,7 @@
     "features.workflows.core",
     "features.release_pipelines.core",
     "segments",
+    "segment_membership",
     "app",
     "e2etests",
     "simple_history",

api/environments/identities/models.py

@@ -5,6 +5,7 @@
 from flag_engine.engine import get_evaluation_result
 
 from environments.identities.managers import IdentityManager
+from environments.identities.signals import traits_changed
 from environments.identities.traits.models import Trait
 from environments.models import Environment
 from environments.sdk.types import SDKTraitData
@@ -211,6 +212,12 @@ def generate_traits(
 
         if persist:
             Trait.objects.bulk_create(trait_models_to_persist)
+            if trait_models_to_persist:
+                traits_changed.send(
+                    sender=type(self),
+                    instance=self,
+                    changed_keys={t.trait_key for t in trait_models_to_persist},
+                )
 
         return trait_models
 
@@ -290,6 +297,18 @@ def update_traits(
         # See: https://github.com/Flagsmith/flagsmith/issues/370
         Trait.objects.bulk_create(new_traits, ignore_conflicts=True)
 
+        changed_keys = (
+            keys_to_delete
+            | {t.trait_key for t in updated_traits}
+            | {t.trait_key for t in new_traits}
+        )
+        if changed_keys:
+            traits_changed.send(
+                sender=type(self),
+                instance=self,
+                changed_keys=changed_keys,
+            )
+
         # return the full list of traits for this identity
         # override persisted traits by transient traits in case of key collisions
         return [

api/environments/identities/signals.py

@@ -0,0 +1,18 @@
+"""Signals emitted from the identity write paths.
+
+`traits_changed` fires after the bulk trait write paths
+(`Identity.update_traits`, `Identity.generate_traits(persist=True)`) complete.
+This is necessary because Django's `post_save` is not emitted by
+`bulk_create` / `bulk_update`, so any consumer that needs to react to trait
+changes from the SDK ingestion path has to subscribe to this signal instead.
+
+Provides:
+
+  * `sender`        — the `Identity` model class.
+  * `instance`      — the identity whose traits changed.
+  * `changed_keys`  — set[str] of trait keys created, updated, or deleted.
+"""
+
+from django.dispatch import Signal
+
+traits_changed = Signal()

api/pyproject.toml

@@ -190,6 +190,7 @@ structlog = "^24.4.0"
 prometheus-client = "^0.21.1"
 django_cockroachdb = "~4.2"
 django-oauth-toolkit = "^3.0.1"
+pyroaring = ">=1.0"
 
 [tool.poetry.group.auth-controller]
 optional = true

api/segment_membership/apps.py

@@ -0,0 +1,15 @@
+from core.apps import BaseAppConfig
+
+
+class SegmentMembershipConfig(BaseAppConfig):
+    name = "segment_membership"
+    default = True
+
+    def ready(self) -> None:
+        super().ready()  # type: ignore[no-untyped-call]
+        # Import order matters — tasks register with the task_processor at
+        # import time, and signals enqueue those tasks.
+        from segment_membership import (
+            signals,  # noqa: F401
+            tasks,  # noqa: F401
+        )

api/segment_membership/constants.py

@@ -0,0 +1,17 @@
+"""Constants for the segment membership index."""
+
+from typing import Literal
+
+# Atom kinds — partition by the property class. Set by
+# `map_segment_condition_to_atom_kind` in `mappers.py`.
+AtomKind = Literal["trait", "identifier", "identity_key", "environment_name"]
+
+KIND_TRAIT: AtomKind = "trait"
+KIND_IDENTIFIER: AtomKind = "identifier"
+KIND_IDENTITY_KEY: AtomKind = "identity_key"
+KIND_ENVIRONMENT_NAME: AtomKind = "environment_name"
+
+# JSONPath properties recognised by the engine for non-trait context values.
+PROPERTY_IDENTITY_IDENTIFIER = "$.identity.identifier"
+PROPERTY_IDENTITY_KEY = "$.identity.key"
+PROPERTY_ENVIRONMENT_NAME = "$.environment.name"

api/segment_membership/dataclasses.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+from segment_membership.constants import AtomKind
+
+
+@dataclass(frozen=True)
+class AtomKey:
+    """Canonical key identifying an atom within an environment."""
+
+    kind: AtomKind
+    property: str
+    operator: str
+    operand_canonical: Optional[str]
+    segment_key: Optional[str]
+
+
+@dataclass(frozen=True)
+class AtomNode:
+    key: AtomKey
+    negated: bool = False
+
+
+@dataclass
+class AndNode:
+    children: list["PredicateTree"] = field(default_factory=list)
+
+
+@dataclass
+class OrNode:
+    children: list["PredicateTree"] = field(default_factory=list)
+
+
+@dataclass
+class TrueNode:
+    """Universe — all ordinals in the env. Lets nested rules with empty
+    conditions short-circuit cleanly."""
+
+
+@dataclass
+class FalseNode:
+    """Empty set — never a member."""
+
+
+PredicateTree = Union[AtomNode, AndNode, OrNode, TrueNode, FalseNode]