feat(analytics): hierarchical groups support

docs/index.rst

@@ -10,6 +10,9 @@ Primary interface
 .. autoclass:: Mixpanel
    :members:
 
+.. autoclass:: MixpanelIdentity
+   :members: is_hierarchical
+
 
 Built-in consumers
 ------------------
@@ -36,3 +39,7 @@ Exceptions
 ----------
 
 .. autoexception:: MixpanelException
+
+.. autoexception:: UnsupportedOperationError
+
+.. autoexception:: CompanyRequiredError

mixpanel/flags/local_feature_flags.py

@@ -3,6 +3,7 @@
 import asyncio
 import copy
 import logging
+import re
 import threading
 import time
 from datetime import datetime, timedelta
@@ -12,6 +13,7 @@
 import json_logic
 
 from mixpanel.credentials import ServiceAccountCredentials
+from mixpanel.identity import MixpanelIdentity
 
 from . import custom_operators  # noqa: F401  (registers custom JsonLogic operators)
 from .types import (
@@ -27,14 +29,44 @@
     REQUEST_HEADERS,
     close_async_client_from_sync,
     dispatch_exposure,
+    exposure_identity,
     generate_traceparent,
+    missing_exposure_identity_message,
     normalized_hash,
     prepare_common_query_params,
 )
 
 logger = logging.getLogger(__name__)
 logging.getLogger("httpx").setLevel(logging.ERROR)
 
+# A {key} placeholder in a context value template.
+_CONTEXT_PLACEHOLDER = re.compile(r"\{([^{}]+)\}")
+
+
+def _apply_context_value_template(
+    template: str | None, context: dict[str, Any]
+) -> str | None:
+    """Fill a flag's context value template from the evaluation context.
+
+    Each ``{key}`` placeholder is replaced with the string value of ``key`` from
+    the context; text outside the braces is literal. Returns None if there is no
+    template or any placeholder's key is missing from the context.
+    """
+    if not template:
+        return None
+    missing = False
+
+    def fill(match: re.Match) -> str:
+        nonlocal missing
+        value = context.get(match.group(1))
+        if value is None:
+            missing = True
+            return ""
+        return str(value)
+
+    resolved = _CONTEXT_PLACEHOLDER.sub(fill, template)
+    return None if missing else resolved
+
 
 class LocalFeatureFlagsProvider:
     FLAGS_DEFINITIONS_URL_PATH = "/flags/definitions"
@@ -46,6 +78,7 @@ def __init__(
         version: str,
         tracker: Callable,
         credentials: ServiceAccountCredentials | None = None,
+        hierarchical_groups: bool = False,
     ) -> None:
         """Initialize the LocalFeatureFlagsProvider.
 
@@ -54,12 +87,17 @@ def __init__(
         :param str version: the version of the Mixpanel library being used, just for tracking
         :param Callable tracker: A function used to track flags exposure events to mixpanel
         :param ServiceAccountCredentials credentials: Optional service account credentials for authentication.
+        :param bool hierarchical_groups: True if the project uses hierarchical groups (default False).
+            Exposure events are then attributed by the context's ``$company_id`` (with ``$user_id``)
+            as well as ``distinct_id``, and the tracker receives a :class:`~mixpanel.identity.MixpanelIdentity`
+            instead of a distinct id string.
         """
         self._token: str = token
         self._config: LocalFlagsConfig = config
         self._version = version
         self._tracker: Callable = tracker
         self._credentials = credentials
+        self._hierarchical_groups = hierarchical_groups
 
         self._flag_definitions: dict[str, ExperimentationFlag] = {}
         self._are_flags_ready = False
@@ -235,7 +273,13 @@ def get_variant(
             logger.warning("Cannot find flag definition for key: '%s'", flag_key)
             return fallback_value.as_fallback(FallbackReason.flag_not_found())
 
-        if not (context_value := context.get(flag_definition.context)):
+        # The flag's template builds the assignment key from the context. Without
+        # one, or with a placeholder missing, the raw context value is used.
+        context_value = _apply_context_value_template(
+            flag_definition.context_value_template, context
+        ) or context.get(flag_definition.context)
+
+        if not context_value:
             logger.warning(
                 "The rollout context, '%s' for flag, '%s' is not present in the supplied context dictionary",
                 flag_definition.context,
@@ -248,7 +292,7 @@ def get_variant(
         selected_variant: SelectedVariant | None = None
 
         if test_user_variant := self._get_variant_override_for_test_user(
-            flag_definition, context
+            flag_definition, context, context_value
         ):
             selected_variant = test_user_variant
         elif rollout := self._get_assigned_rollout(
@@ -288,13 +332,22 @@ def track_exposure_event(
         self._track_exposure(flag_key, variant, context)
 
     def _get_variant_override_for_test_user(
-        self, flag_definition: ExperimentationFlag, context: dict[str, Any]
+        self,
+        flag_definition: ExperimentationFlag,
+        context: dict[str, Any],
+        context_value: Any,
     ) -> SelectedVariant | None:
         """Check if user has a test variant override."""
         if not flag_definition.ruleset.test or not flag_definition.ruleset.test.users:
             return None
 
-        if not (distinct_id := context.get("distinct_id")):
+        # Overrides are keyed by distinct id. A flag assigned on distinct_id
+        # looks them up by its resolved assignment key.
+        if flag_definition.context == "distinct_id":
+            distinct_id = context_value
+        else:
+            distinct_id = context.get("distinct_id")
+        if not distinct_id:
             return None
 
         if not (variant_key := flag_definition.ruleset.test.users.get(distinct_id)):
@@ -522,7 +575,9 @@ def _track_exposure(
         context: dict[str, Any],
         latency_in_seconds: float | None = None,
     ):
-        if distinct_id := context.get("distinct_id"):
+        if (
+            identity := exposure_identity(context, self._hierarchical_groups)
+        ) is not None:
             properties = {
                 "Experiment name": flag_key,
                 "Variant name": variant.variant_key,
@@ -536,13 +591,13 @@ def _track_exposure(
             if latency_in_seconds is not None:
                 properties["Variant fetch latency (ms)"] = latency_in_seconds * 1000
 
-            self._dispatch_exposure(distinct_id, properties)
+            self._dispatch_exposure(identity, properties)
         else:
-            logger.error(
-                "Cannot track exposure event without a distinct_id in the context"
-            )
+            logger.error(missing_exposure_identity_message(self._hierarchical_groups))
 
-    def _dispatch_exposure(self, distinct_id: str, properties: dict[str, Any]) -> None:
+    def _dispatch_exposure(
+        self, distinct_id: str | MixpanelIdentity, properties: dict[str, Any]
+    ) -> None:
         dispatch_exposure(
             self._tracker, self._config.exposure_executor, distinct_id, properties
         )

mixpanel/flags/remote_feature_flags.py

@@ -10,6 +10,7 @@
 from asgiref.sync import sync_to_async
 
 from mixpanel.credentials import ServiceAccountCredentials
+from mixpanel.identity import MixpanelIdentity
 
 from .types import (
     FallbackReason,
@@ -23,7 +24,9 @@
     REQUEST_HEADERS,
     close_async_client_from_sync,
     dispatch_exposure,
+    exposure_identity,
     generate_traceparent,
+    missing_exposure_identity_message,
     prepare_common_query_params,
 )
 
@@ -41,6 +44,7 @@ def __init__(
         version: str,
         tracker: Callable,
         credentials: ServiceAccountCredentials | None = None,
+        hierarchical_groups: bool = False,
     ) -> None:
         """Initialize the RemoteFeatureFlagsProvider.
 
@@ -49,12 +53,17 @@ def __init__(
         :param str version: the version of the Mixpanel library being used, just for tracking
         :param Callable tracker: A function used to track flags exposure events to mixpanel
         :param ServiceAccountCredentials credentials: Optional service account credentials for authentication.
+        :param bool hierarchical_groups: True if the project uses hierarchical groups (default False).
+            Exposure events are then attributed by the context's ``$company_id`` (with ``$user_id``)
+            as well as ``distinct_id``, and the tracker receives a :class:`~mixpanel.identity.MixpanelIdentity`
+            instead of a distinct id string.
         """
         self._token: str = token
         self._config: RemoteFlagsConfig = config
         self._version: str = version
         self._tracker: Callable = tracker
         self._credentials = credentials
+        self._hierarchical_groups = hierarchical_groups
         self._project_id: str | None = credentials.project_id if credentials else None
 
         # Build httpx client parameters
@@ -151,14 +160,15 @@ async def aget_variant(
             if (
                 not is_fallback
                 and reportExposure
-                and (distinct_id := context.get("distinct_id"))
+                and (identity := exposure_identity(context, self._hierarchical_groups))
+                is not None
             ):
                 properties = self._build_tracking_properties(
                     flag_key, selected_variant, start_time, end_time
                 )
                 asyncio.create_task(  # noqa: RUF006 - intentional fire-and-forget for exposure tracking
                     sync_to_async(self._tracker, thread_sensitive=False)(
-                        distinct_id, EXPOSURE_EVENT, properties
+                        identity, EXPOSURE_EVENT, properties
                     )
                 )
         except Exception as exc:
@@ -193,16 +203,16 @@ async def atrack_exposure_event(
         :param SelectedVariant variant: The selected variant for the feature flag
         :param Dict[str, Any] context: The user context used to evaluate the feature flag
         """
-        if distinct_id := context.get("distinct_id"):
+        if (
+            identity := exposure_identity(context, self._hierarchical_groups)
+        ) is not None:
             properties = self._build_tracking_properties(flag_key, variant)
 
             await sync_to_async(self._tracker, thread_sensitive=False)(
-                distinct_id, EXPOSURE_EVENT, properties
+                identity, EXPOSURE_EVENT, properties
             )
         else:
-            logger.error(
-                "Cannot track exposure event without a distinct_id in the context"
-            )
+            logger.error(missing_exposure_identity_message(self._hierarchical_groups))
 
     def get_all_variants(
         self, context: dict[str, Any]
@@ -276,12 +286,13 @@ def get_variant(
             if (
                 not is_fallback
                 and reportExposure
-                and (distinct_id := context.get("distinct_id"))
+                and (identity := exposure_identity(context, self._hierarchical_groups))
+                is not None
             ):
                 properties = self._build_tracking_properties(
                     flag_key, selected_variant, start_time, end_time
                 )
-                self._dispatch_exposure(distinct_id, properties)
+                self._dispatch_exposure(identity, properties)
 
         except Exception as exc:
             logger.exception("Failed to get remote variant for flag '%s'", flag_key)
@@ -313,15 +324,17 @@ def track_exposure_event(
         :param SelectedVariant variant: The selected variant for the feature flag
         :param Dict[str, Any] context: The user context used to evaluate the feature flag
         """
-        if distinct_id := context.get("distinct_id"):
+        if (
+            identity := exposure_identity(context, self._hierarchical_groups)
+        ) is not None:
             properties = self._build_tracking_properties(flag_key, variant)
-            self._dispatch_exposure(distinct_id, properties)
+            self._dispatch_exposure(identity, properties)
         else:
-            logger.error(
-                "Cannot track exposure event without a distinct_id in the context"
-            )
+            logger.error(missing_exposure_identity_message(self._hierarchical_groups))
 
-    def _dispatch_exposure(self, distinct_id: str, properties: dict[str, Any]) -> None:
+    def _dispatch_exposure(
+        self, distinct_id: str | MixpanelIdentity, properties: dict[str, Any]
+    ) -> None:
         dispatch_exposure(
             self._tracker, self._config.exposure_executor, distinct_id, properties
         )