from opentelemetry.context import Context
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, ConsoleSpanExporter
from opentelemetry.sdk.trace.sampling import Sampler, SamplingResult, Decision, ParentBasedTraceIdRatio
from opentelemetry import trace
from opentelemetry.util.types import Attributes # For type hinting
USER_ID_TO_DROP = "user_to_skip_tracing"
class UserBasedSampler(Sampler):
# A custom sampler that drops any span having a `user.id` attribute matching
# a specified user ID. For other cases, it delegates to a root sampler.
def __init__(self, root_sampler: Sampler = ParentBasedTraceIdRatio(0.5)):
self._root_sampler = root_sampler
def should_sample(
self,
parent_context: Context,
trace_id: int,
name: str,
kind, # SpanKind is implicitly an int here
attributes: Attributes,
links
) -> SamplingResult:
user_id = attributes.get("user.id") if attributes else None
if user_id == USER_ID_TO_DROP:
return SamplingResult(
decision=Decision.DROP,
attributes={"sampler.reason": f"Dropping span for user.id={user_id}"}
)
else:
return self._root_sampler.should_sample(parent_context, trace_id, name, kind, attributes, links)
def get_description(self) -> str:
return f"UserBasedSampler(root_sampler={self._root_sampler.get_description()})"
# Example usage:
# if __name__ == "__main__":
# custom_sampler = UserBasedSampler(root_sampler=ParentBasedTraceIdRatio(1.0))
# provider = TracerProvider(sampler=custom_sampler)
# trace.set_tracer_provider(provider)
# provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
# tracer = trace.get_tracer(__name__, "0.1.0")
# with tracer.start_as_current_span("op_for_dropped_user", attributes={"user.id": USER_ID_TO_DROP}): pass
# with tracer.start_as_current_span("op_for_sampled_user", attributes={"user.id": "another_user"}): pass
# with tracer.start_as_current_span("op_without_user_id"): pass