feat: Propagate environment ID to evaluation hooks

Co-Authored-By: rlamb@launchdarkly.com <4955475+kinyoklion@users.noreply.github.com>

Author: devin-ai-integration[bot]

ldclient/client.py

@@ -709,7 +709,7 @@ def __evaluate_with_hooks(self, key: str, context: Context, default_value: Any,
 
             hooks = self.__hooks.copy()
 
-        series_context = EvaluationSeriesContext(key=key, context=context, default_value=default_value, method=method)
+        series_context = EvaluationSeriesContext(key=key, context=context, default_value=default_value, method=method, environment_id=self._data_system.environment_id)
         hook_data = self.__execute_before_evaluation(hooks, series_context)
         evaluation_result = block()
         self.__execute_after_evaluation(hooks, series_context, hook_data, evaluation_result.evaluation_detail)

ldclient/hook.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 
 from ldclient.context import Context
 from ldclient.evaluation import EvaluationDetail
@@ -17,6 +17,7 @@ class EvaluationSeriesContext:
     context: Context  #: The context used during evaluation.
     default_value: Any  #: The default value provided to the evaluation method
     method: str  #: The string version of the method which triggered the evaluation series.
+    environment_id: Optional[str] = None  #: The environment ID the SDK is connected to, if available.
 
 
 @dataclass

ldclient/impl/datasource/datasource_common.py

@@ -5,8 +5,9 @@
 # currently excluded from documentation - see docs/README.md
 
 from collections import namedtuple
-from typing import Optional
+from typing import Mapping, Optional, Protocol, runtime_checkable
 
+from ldclient.impl.util import _LD_ENVID_HEADER
 from ldclient.interfaces import DataSourceUpdateSink, FeatureStore
 from ldclient.versioned_data_kind import FEATURES, SEGMENTS
 
@@ -34,6 +35,32 @@ def sink_or_store(sink: Optional[DataSourceUpdateSink], store: FeatureStore):
     return sink
 
 
+@runtime_checkable
+class EnvironmentIdSink(Protocol):
+    """
+    Implemented by data source update sinks which can record the environment ID
+    reported by LaunchDarkly. This is separate from
+    :class:`ldclient.interfaces.DataSourceUpdateSink` so that externally
+    implemented sinks remain compatible.
+    """
+
+    def set_environment_id(self, environment_id: str) -> None:
+        ...
+
+
+def record_environment_id(sink, headers: Optional[Mapping[str, str]]):
+    """
+    Records the environment ID from a set of LaunchDarkly response headers, if
+    both the headers and the sink provide one.
+    """
+    if headers is None or not isinstance(sink, EnvironmentIdSink):
+        return
+
+    environment_id = headers.get(_LD_ENVID_HEADER)
+    if isinstance(environment_id, str) and environment_id != '':
+        sink.set_environment_id(environment_id)
+
+
 def parse_path(path: str):
     for kind in [FEATURES, SEGMENTS]:
         if path.startswith(kind.stream_api_path):

ldclient/impl/datasource/feature_requester.py

@@ -4,6 +4,7 @@
 
 import json
 from collections import namedtuple
+from typing import Mapping, Optional, Tuple
 from urllib import parse
 
 import urllib3
@@ -27,6 +28,10 @@ def __init__(self, config):
             self._poll_uri += '?%s' % parse.urlencode({'filter': config.payload_filter_key})
 
     def get_all_data(self):
+        (data, _) = self.get_all_data_with_headers()
+        return data
+
+    def get_all_data_with_headers(self) -> Tuple[dict, Optional[Mapping[str, str]]]:
         uri = self._poll_uri
         hdrs = _headers(self._config)
         cache_entry = self._cache.get(uri)
@@ -47,4 +52,4 @@ def get_all_data(self):
                 self._cache[uri] = CacheEntry(data=data, etag=etag)
         log.debug("%s response status:[%d] From cache? [%s] ETag:[%s]", uri, r.status, from_cache, etag)
 
-        return {FEATURES: data['flags'], SEGMENTS: data['segments']}
+        return ({FEATURES: data['flags'], SEGMENTS: data['segments']}, r.headers)

ldclient/impl/datasource/polling.py

@@ -6,10 +6,13 @@
 
 import time
 from threading import Event
-from typing import Optional
+from typing import Any, Mapping, Optional, Protocol, Tuple, runtime_checkable
 
 from ldclient.config import Config
-from ldclient.impl.datasource.datasource_common import sink_or_store
+from ldclient.impl.datasource.datasource_common import (
+    record_environment_id,
+    sink_or_store
+)
 from ldclient.impl.repeating_task import RepeatingTask
 from ldclient.impl.util import (
     UnsuccessfulResponseException,
@@ -28,6 +31,12 @@
 )
 
 
+@runtime_checkable
+class _FeatureRequesterWithHeaders(Protocol):
+    def get_all_data_with_headers(self) -> Tuple[Any, Optional[Mapping[str, str]]]:
+        ...
+
+
 class PollingUpdateProcessor(UpdateProcessor):
     def __init__(self, config: Config, requester: FeatureRequester, store: FeatureStore, ready: Event):
         self._config = config
@@ -58,7 +67,8 @@ def __stop_with_error_info(self, error: Optional[DataSourceErrorInfo]):
 
     def _poll(self):
         try:
-            all_data = self._requester.get_all_data()
+            (all_data, headers) = self._get_all_data_with_headers()
+            record_environment_id(self._data_source_update_sink, headers)
             sink_or_store(self._data_source_update_sink, self._store).init(all_data)
             if not self._ready.is_set() and self._store.initialized:
                 log.info("PollingUpdateProcessor initialized ok")
@@ -84,3 +94,13 @@ def _poll(self):
 
             if self._data_source_update_sink is not None:
                 self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, DataSourceErrorInfo(DataSourceErrorKind.UNKNOWN, 0, time.time(), str(e)))
+
+    def _get_all_data_with_headers(self) -> Tuple[Any, Optional[Mapping[str, str]]]:
+        """
+        Externally provided feature requesters are not required to surface
+        response headers, so fall back to the data-only method.
+        """
+        if isinstance(self._requester, _FeatureRequesterWithHeaders):
+            return self._requester.get_all_data_with_headers()
+
+        return (self._requester.get_all_data(), None)

ldclient/impl/datasource/status.py

@@ -26,12 +26,22 @@ def __init__(self, store: FeatureStore, status_listeners: Listeners, flag_change
 
         self.__lock = ReadWriteLock()
         self.__status = DataSourceStatus(DataSourceState.INITIALIZING, time.time(), None)
+        self.__environment_id: Optional[str] = None
 
     @property
     def status(self) -> DataSourceStatus:
         with self.__lock.read():
             return self.__status
 
+    @property
+    def environment_id(self) -> Optional[str]:
+        with self.__lock.read():
+            return self.__environment_id
+
+    def set_environment_id(self, environment_id: str) -> None:
+        with self.__lock.write():
+            self.__environment_id = environment_id
+
     def init(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]):
         old_data = None
 

ldclient/impl/datasource/streaming.py

@@ -5,7 +5,7 @@
 from urllib import parse
 
 from ld_eventsource import SSEClient
-from ld_eventsource.actions import Event, Fault
+from ld_eventsource.actions import Event, Fault, Start
 from ld_eventsource.config import (
     ConnectStrategy,
     ErrorStrategy,
@@ -16,6 +16,7 @@
 from ldclient.impl.datasource.datasource_common import (
     STREAM_ALL_PATH,
     parse_path,
+    record_environment_id,
     sink_or_store
 )
 from ldclient.impl.http import HTTPFactory, _http_factory
@@ -62,7 +63,9 @@ def run(self):
         self._sse = self._create_sse_client()
         self._connection_attempt_start_time = time.time()
         for action in self._sse.all:
-            if isinstance(action, Event):
+            if isinstance(action, Start):
+                record_environment_id(self._data_source_update_sink, action.headers)
+            elif isinstance(action, Event):
                 message_ok = False
                 try:
                     message_ok = self._process_message(sink_or_store(self._data_source_update_sink, self._store), action)
@@ -91,6 +94,8 @@ def run(self):
                         log.info("StreamingUpdateProcessor initialized ok.")
                         self._ready.set()
             elif isinstance(action, Fault):
+                record_environment_id(self._data_source_update_sink, action.headers)
+
                 # If the SSE client detects the stream has closed, then it will emit a fault with no-error. We can
                 # ignore this since we want the connection to continue.
                 if action.error is None:

ldclient/impl/datasystem/__init__.py

@@ -6,7 +6,7 @@
 from abc import abstractmethod
 from enum import Enum
 from threading import Event
-from typing import Protocol, runtime_checkable
+from typing import Optional, Protocol, runtime_checkable
 
 from ldclient.impl.listeners import Listeners
 from ldclient.interfaces import (
@@ -142,6 +142,18 @@ def store(self) -> ReadOnlyStore:
         """
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def environment_id(self) -> Optional[str]:
+        """
+        Returns the environment ID reported by LaunchDarkly, if it is known.
+
+        This is only available once a connection to LaunchDarkly has provided
+        it, and it will be None when the SDK is offline, using an unsupported
+        data source, or connected to a service which does not report it.
+        """
+        raise NotImplementedError
+
 
 class DiagnosticAccumulator(Protocol):
     def record_stream_init(self, timestamp, duration, failed):

ldclient/impl/datasystem/fdv1.py

@@ -98,6 +98,10 @@ def stop(self):
     def store(self) -> ReadOnlyStore:
         return self._store_wrapper
 
+    @property
+    def environment_id(self) -> Optional[str]:
+        return self._data_source_update_sink.environment_id
+
     def set_diagnostic_accumulator(self, diagnostic_accumulator: DiagnosticAccumulator):
         """
         Sets the diagnostic accumulator for streaming initialization metrics.

ldclient/impl/datasystem/fdv2.py

@@ -21,7 +21,12 @@
 from ldclient.impl.listeners import Listeners
 from ldclient.impl.repeating_task import RepeatingTask
 from ldclient.impl.rwlock import ReadWriteLock
-from ldclient.impl.util import _LD_FD_FALLBACK_HEADER, _Fail, log
+from ldclient.impl.util import (
+    _LD_ENVID_HEADER,
+    _LD_FD_FALLBACK_HEADER,
+    _Fail,
+    log
+)
 from ldclient.interfaces import (
     DataSourceErrorInfo,
     DataSourceErrorKind,
@@ -95,6 +100,7 @@ def __init__(
         self._lock = ReadWriteLock()
         self._active_synchronizer: Optional[Synchronizer] = None
         self._threads: List[Thread] = []
+        self._environment_id: Optional[str] = None
 
         # Track configuration
         self._configured_with_data_sources = (
@@ -216,6 +222,8 @@ def _run_initializers(self, set_on_ready: Event) -> bool:
 
                 if isinstance(basis_result, _Fail):
                     log.warning("Initializer %s failed: %s", initializer.name, basis_result.error)
+                    if basis_result.headers is not None:
+                        self._record_environment_id(basis_result.headers.get(_LD_ENVID_HEADER))
                     # An error response can still carry the FDv1 fallback directive.
                     if basis_result.headers is not None and \
                             basis_result.headers.get(_LD_FD_FALLBACK_HEADER) == 'true':
@@ -240,6 +248,8 @@ def _run_initializers(self, set_on_ready: Event) -> bool:
                 basis = basis_result.value
                 log.info("Initialized via %s", initializer.name)
 
+                self._record_environment_id(basis.environment_id)
+
                 # Apply the basis to the store
                 self._store.apply(basis.change_set, basis.persist)
 
@@ -412,6 +422,8 @@ def reader(self: 'FDv2'):
                 if self._stop_event.is_set():
                     return ConditionDirective.FALLBACK
 
+                self._record_environment_id(update.environment_id)
+
                 # Handle the update
                 if update.change_set is not None:
                     self._store.apply(update.change_set, True)
@@ -496,6 +508,19 @@ def _persistent_store_outage_recovery(self, data_store_status: DataStoreStatus):
         if err is not None:
             log.error("Failed to reinitialize data store", exc_info=err)
 
+    def _record_environment_id(self, environment_id: Optional[str]):
+        if environment_id is None:
+            return
+
+        with self._lock.write():
+            self._environment_id = environment_id
+
+    @property
+    def environment_id(self) -> Optional[str]:
+        """Get the environment ID reported by LaunchDarkly, if known."""
+        with self._lock.read():
+            return self._environment_id
+
     @property
     def store(self) -> ReadOnlyStore:
         """Get the underlying store for flag evaluation."""