Merge pull request #26 from UiPath/fix/add_trace_settings

fix: add trace settings to filter spans before export

Author: cristipufu

pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "uipath-core"
-version = "0.1.10"
+version = "0.2.0"
 description = "UiPath Core abstractions"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"
 dependencies = [
-  "opentelemetry-sdk>=1.39.0, <2.0.0",
-  "opentelemetry-instrumentation>=0.60b0, <1.0.0",
+  "opentelemetry-sdk>=1.39.1, <2.0.0",
+  "opentelemetry-instrumentation>=0.60b1, <1.0.0",
   "pydantic>=2.12.5, <3.0.0",
 ]
 classifiers = [

src/uipath/core/tracing/__init__.py

@@ -7,9 +7,11 @@
 from uipath.core.tracing.decorators import traced
 from uipath.core.tracing.span_utils import UiPathSpanUtils
 from uipath.core.tracing.trace_manager import UiPathTraceManager
+from uipath.core.tracing.types import UiPathTraceSettings
 
 __all__ = [
     "traced",
     "UiPathSpanUtils",
     "UiPathTraceManager",
+    "UiPathTraceSettings",
 ]

src/uipath/core/tracing/processors.py

@@ -1,22 +1,27 @@
 """Custom span processors for UiPath execution tracing."""
 
-from typing import Optional, cast
+from typing import cast
 
 from opentelemetry import context as context_api
 from opentelemetry import trace
-from opentelemetry.sdk.trace import Span
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.sdk.trace.export import (
     BatchSpanProcessor,
     SimpleSpanProcessor,
+    SpanExporter,
 )
 
+from uipath.core.tracing.types import UiPathTraceSettings
+
 
 class UiPathExecutionTraceProcessorMixin:
-    def on_start(
-        self, span: Span, parent_context: Optional[context_api.Context] = None
-    ):
+    """Mixin that propagates execution.id and optionally filters spans."""
+
+    _settings: UiPathTraceSettings | None = None
+
+    def on_start(self, span: Span, parent_context: context_api.Context | None = None):
         """Called when a span is started."""
-        parent_span: Optional[Span]
+        parent_span: Span | None
         if parent_context:
             parent_span = cast(Span, trace.get_current_span(parent_context))
         else:
@@ -27,17 +32,42 @@ def on_start(
             if execution_id:
                 span.set_attribute("execution.id", execution_id)
 
+    def on_end(self, span: ReadableSpan):
+        """Called when a span ends. Filters before delegating to parent."""
+        span_filter = self._settings.span_filter if self._settings else None
+        if span_filter is None or span_filter(span):
+            parent = cast(SpanProcessor, super())
+            parent.on_end(span)
+
 
 class UiPathExecutionBatchTraceProcessor(
     UiPathExecutionTraceProcessorMixin, BatchSpanProcessor
 ):
-    """Batch span processor that propagates execution.id."""
+    """Batch span processor that propagates execution.id and optionally filters."""
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        settings: UiPathTraceSettings | None = None,
+    ):
+        """Initialize the batch trace processor."""
+        super().__init__(span_exporter)
+        self._settings = settings
 
 
 class UiPathExecutionSimpleTraceProcessor(
     UiPathExecutionTraceProcessorMixin, SimpleSpanProcessor
 ):
-    """Simple span processor that propagates execution.id."""
+    """Simple span processor that propagates execution.id and optionally filters."""
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        settings: UiPathTraceSettings | None = None,
+    ):
+        """Initialize the simple trace processor."""
+        super().__init__(span_exporter)
+        self._settings = settings
 
 
 __all__ = [

src/uipath/core/tracing/trace_manager.py

@@ -15,6 +15,7 @@
     UiPathExecutionBatchTraceProcessor,
     UiPathExecutionSimpleTraceProcessor,
 )
+from uipath.core.tracing.types import UiPathTraceSettings
 
 
 class UiPathTraceManager:
@@ -40,13 +41,22 @@ def add_span_exporter(
         self,
         span_exporter: SpanExporter,
         batch: bool = True,
+        settings: UiPathTraceSettings | None = None,
     ) -> UiPathTraceManager:
-        """Add a span processor to the tracer provider."""
+        """Add a span exporter to the tracer provider.
+
+        Args:
+            span_exporter: The exporter to add.
+            batch: Whether to use batch processing (default: True).
+            settings: Optional trace settings for filtering, etc.
+        """
         span_processor: SpanProcessor
         if batch:
-            span_processor = UiPathExecutionBatchTraceProcessor(span_exporter)
+            span_processor = UiPathExecutionBatchTraceProcessor(span_exporter, settings)
         else:
-            span_processor = UiPathExecutionSimpleTraceProcessor(span_exporter)
+            span_processor = UiPathExecutionSimpleTraceProcessor(
+                span_exporter, settings
+            )
         self.tracer_span_processors.append(span_processor)
         self.tracer_provider.add_span_processor(span_processor)
         return self

src/uipath/core/tracing/types.py

@@ -0,0 +1,21 @@
+"""Tracing types for UiPath SDK."""
+
+from typing import Callable
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field
+
+
+class UiPathTraceSettings(BaseModel):
+    """Trace settings for UiPath SDK."""
+
+    model_config = {"arbitrary_types_allowed": True}  # Needed for Callable
+
+    span_filter: Callable[[ReadableSpan], bool] | None = Field(
+        default=None,
+        description=(
+            "Optional filter to decide whether a span should be exported. "
+            "Called when a span ends with a ReadableSpan argument. "
+            "Return True to export, False to skip."
+        ),
+    )