Refactoring propagation context by antonpirker · Pull Request #2970 · getsentry/sentry-python

@@ -1,6 +1,5 @@ import os import sys import uuid from copy import copy from collections import deque from contextlib import contextmanager Expand All @@ -15,9 +14,9 @@ from sentry_sdk.session import Session from sentry_sdk.tracing_utils import ( Baggage, extract_sentrytrace_data, has_tracing_enabled, normalize_incoming_data, PropagationContext, ) from sentry_sdk.tracing import ( BAGGAGE_HEADER_NAME, Expand Down Expand Up @@ -196,7 +195,7 @@ def __init__(self, ty=None, client=None): self._error_processors = [] # type: List[ErrorProcessor]
self._name = None # type: Optional[str] self._propagation_context = None # type: Optional[Dict[str, Any]] self._propagation_context = None # type: Optional[PropagationContext]
self.client = NonRecordingClient() # type: sentry_sdk.client.BaseClient
Expand Down Expand Up @@ -431,77 +430,28 @@ def _load_trace_data_from_env(self):
return incoming_trace_information or None
def _extract_propagation_context(self, data): # type: (Dict[str, Any]) -> Optional[Dict[str, Any]] context = {} # type: Dict[str, Any] normalized_data = normalize_incoming_data(data)
baggage_header = normalized_data.get(BAGGAGE_HEADER_NAME) if baggage_header: context["dynamic_sampling_context"] = Baggage.from_incoming_header( baggage_header ).dynamic_sampling_context()
sentry_trace_header = normalized_data.get(SENTRY_TRACE_HEADER_NAME) if sentry_trace_header: sentrytrace_data = extract_sentrytrace_data(sentry_trace_header) if sentrytrace_data is not None: context.update(sentrytrace_data)
only_baggage_no_sentry_trace = ( "dynamic_sampling_context" in context and "trace_id" not in context ) if only_baggage_no_sentry_trace: context.update(self._create_new_propagation_context())
if context: if not context.get("span_id"): context["span_id"] = uuid.uuid4().hex[16:]
return context
return None
def _create_new_propagation_context(self): # type: () -> Dict[str, Any] return { "trace_id": uuid.uuid4().hex, "span_id": uuid.uuid4().hex[16:], "parent_span_id": None, "dynamic_sampling_context": None, }
def set_new_propagation_context(self): # type: () -> None """ Creates a new propagation context and sets it as `_propagation_context`. Overwriting existing one. """ self._propagation_context = self._create_new_propagation_context() logger.debug( "[Tracing] Create new propagation context: %s", self._propagation_context, ) self._propagation_context = PropagationContext()
def generate_propagation_context(self, incoming_data=None): # type: (Optional[Dict[str, str]]) -> None """ Makes sure the propagation context (`_propagation_context`) is set. The propagation context only lives on the current scope. If there is `incoming_data` overwrite existing `_propagation_context`. if there is no `incoming_data` create new `_propagation_context`, but do NOT overwrite if already existing. Makes sure the propagation context is set on the scope. If there is `incoming_data` overwrite existing propagation context. If there is no `incoming_data` create new propagation context, but do NOT overwrite if already existing. """ if incoming_data: context = self._extract_propagation_context(incoming_data)
if context is not None: self._propagation_context = context logger.debug( "[Tracing] Extracted propagation context from incoming data: %s", self._propagation_context, ) propagation_context = PropagationContext.from_incoming_data(incoming_data) if propagation_context is not None: self._propagation_context = propagation_context
if self._propagation_context is None and self._type != ScopeType.CURRENT: self.set_new_propagation_context() if self._type != ScopeType.CURRENT: if self._propagation_context is None: self.set_new_propagation_context()
def get_dynamic_sampling_context(self): # type: () -> Optional[Dict[str, str]] Expand All @@ -514,11 +464,11 @@ def get_dynamic_sampling_context(self):
baggage = self.get_baggage() if baggage is not None: self._propagation_context["dynamic_sampling_context"] = ( self._propagation_context.dynamic_sampling_context = ( baggage.dynamic_sampling_context() )
return self._propagation_context["dynamic_sampling_context"] return self._propagation_context.dynamic_sampling_context
def get_traceparent(self, *args, **kwargs): # type: (Any, Any) -> Optional[str] Expand All @@ -535,8 +485,8 @@ def get_traceparent(self, *args, **kwargs): # If this scope has a propagation context, return traceparent from there if self._propagation_context is not None: traceparent = "%s-%s" % ( self._propagation_context["trace_id"], self._propagation_context["span_id"], self._propagation_context.trace_id, self._propagation_context.span_id, ) return traceparent
Expand All @@ -557,8 +507,8 @@ def get_baggage(self, *args, **kwargs):
# If this scope has a propagation context, return baggage from there if self._propagation_context is not None: dynamic_sampling_context = self._propagation_context.get( "dynamic_sampling_context" dynamic_sampling_context = ( self._propagation_context.dynamic_sampling_context ) if dynamic_sampling_context is None: return Baggage.from_options(self) Expand All @@ -577,9 +527,9 @@ def get_trace_context(self): return None
trace_context = { "trace_id": self._propagation_context["trace_id"], "span_id": self._propagation_context["span_id"], "parent_span_id": self._propagation_context["parent_span_id"], "trace_id": self._propagation_context.trace_id, "span_id": self._propagation_context.span_id, "parent_span_id": self._propagation_context.parent_span_id, "dynamic_sampling_context": self.get_dynamic_sampling_context(), } # type: Dict[str, Any]
Expand Down Expand Up @@ -667,7 +617,7 @@ def iter_trace_propagation_headers(self, *args, **kwargs): yield header
def get_active_propagation_context(self): # type: () -> Dict[str, Any] # type: () -> Optional[PropagationContext] if self._propagation_context is not None: return self._propagation_context
Expand All @@ -679,7 +629,7 @@ def get_active_propagation_context(self): if isolation_scope._propagation_context is not None: return isolation_scope._propagation_context
return {} return None
def clear(self): # type: () -> None Expand Down Expand Up @@ -1069,12 +1019,11 @@ def start_span(self, instrumenter=INSTRUMENTER.SENTRY, **kwargs): span = self.span or Scope.get_isolation_scope().span
if span is None: # New spans get the `trace_id`` from the scope # New spans get the `trace_id` from the scope if "trace_id" not in kwargs:
trace_id = self.get_active_propagation_context().get("trace_id") if trace_id is not None: kwargs["trace_id"] = trace_id propagation_context = self.get_active_propagation_context() if propagation_context is not None: kwargs["trace_id"] = propagation_context.trace_id
span = Span(**kwargs) else: Expand Down