Correctly trace OTLP spans (#457)

This PR fixes an issue where `DBOS.span` does not return the current
active span. Instead, it always returned the span created by DBOS. This
would cause incorrect span tree structure.

For example:
```python
@DBOS.workflow()
def test_workflow() -> None:
    with my_tracer.start_as_current_span("manual_span"):
        test_step()
```

Without this PR, the span tree is wrong:
```
test_workflow
|- manual_span
|- test_step
```

With this PR, `test_step` will correctly show as the child of
`manual_span`:
```
test_workflow
|- manual_span
    |- test_step
```

Author: qianl15

dbos/_context.py

@@ -215,11 +215,18 @@ def start_handler(self, attributes: TracedAttributes) -> None:
     def end_handler(self, exc_value: Optional[BaseException]) -> None:
         self._end_span(exc_value)
 
-    def get_current_span(self) -> Optional[Span]:
+    """ Return the current DBOS span if any. It must be a span created by DBOS."""
+
+    def get_current_dbos_span(self) -> Optional[Span]:
         if len(self.context_spans) > 0:
             return self.context_spans[-1].span
         return None
 
+    """ Return the current active span if any. It might not be a DBOS span."""
+
+    def get_current_active_span(self) -> Optional[Span]:
+        return dbos_tracer.get_current_span()
+
     def _start_span(self, attributes: TracedAttributes) -> None:
         if dbos_tracer.disable_otlp:
             return
@@ -235,7 +242,7 @@ def _start_span(self, attributes: TracedAttributes) -> None:
         attributes["authenticatedUserAssumedRole"] = self.assumed_role
         span = dbos_tracer.start_span(
             attributes,
-            parent=self.context_spans[-1].span if len(self.context_spans) > 0 else None,
+            parent=None,  # It'll use the current active span as the parent
         )
         # Activate the current span
         cm = use_span(

dbos/_core.py

@@ -971,7 +971,7 @@ def invoke_tx(*args: Any, **kwargs: Any) -> Any:
                                 dbapi_error
                             ) or dbos._app_db._is_serialization_error(dbapi_error):
                                 # Retry on serialization failure
-                                span = ctx.get_current_span()
+                                span = ctx.get_current_dbos_span()
                                 if span:
                                     span.add_event(
                                         "Transaction Failure",
@@ -1090,7 +1090,7 @@ def on_exception(attempt: int, error: BaseException) -> float:
                     exc_info=error,
                 )
                 ctx = assert_current_dbos_context()
-                span = ctx.get_current_span()
+                span = ctx.get_current_dbos_span()
                 if span:
                     span.add_event(
                         f"Step attempt {attempt} failed",

dbos/_dbos.py

@@ -1297,7 +1297,7 @@ def parent_workflow_id(cls) -> str:
     def span(cls) -> Span:
         """Return the tracing `Span` associated with the current context."""
         ctx = assert_current_dbos_context()
-        span = ctx.get_current_span()
+        span = ctx.get_current_active_span()
         assert span
         return span
 

dbos/_logger.py

@@ -39,7 +39,7 @@ def filter(self, record: Any) -> bool:
         if ctx:
             if ctx.is_within_workflow():
                 record.operationUUID = ctx.workflow_id
-            span = ctx.get_current_span()
+            span = ctx.get_current_active_span()
             if span:
                 trace_id = format_trace_id(span.get_span_context().trace_id)
                 record.traceId = trace_id

dbos/_tracer.py

@@ -77,5 +77,12 @@ def start_span(
     def end_span(self, span: Span) -> None:
         span.end()
 
+    def get_current_span(self) -> Optional[Span]:
+        # Return the current active span if any. It might not be a DBOS span.
+        span = trace.get_current_span()
+        if span.get_span_context().is_valid:
+            return span
+        return None
+
 
 dbos_tracer = DBOSTracer()

pdm.lock

@@ -85,4 +85,5 @@ dev = [
     "pyright>=1.1.398",
     "types-docker>=7.1.0.20241229",
     "pytest-timeout>=2.3.1",
+    "inline-snapshot>=0.28.0",
 ]