dbos 1.14.0a4__tar.gz → 1.14.0a5__tar.gz

{dbos-1.14.0a4 → dbos-1.14.0a5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 1.14.0a4
+Version: 1.14.0a5
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-1.14.0a4 → dbos-1.14.0a5}/dbos/_context.py

@@ -215,11 +215,18 @@ class DBOSContext:
     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 @@ class DBOSContext:
         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-1.14.0a4 → dbos-1.14.0a5}/dbos/_core.py

@@ -971,7 +971,7 @@ def decorate_transaction(
                                 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 decorate_step(
                     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-1.14.0a4 → dbos-1.14.0a5}/dbos/_dbos.py

@@ -1297,7 +1297,7 @@ class DBOS:
     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-1.14.0a4 → dbos-1.14.0a5}/dbos/_logger.py

@@ -39,7 +39,7 @@ class DBOSLogTransformer(logging.Filter):
         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-1.14.0a4 → dbos-1.14.0a5}/dbos/_tracer.py

@@ -77,5 +77,12 @@ class DBOSTracer:
     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()

{dbos-1.14.0a4 → dbos-1.14.0a5}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "1.14.0a4"
+version = "1.14.0a5"
 
 [project.license]
 text = "MIT"
@@ -90,4 +90,5 @@ dev = [
     "pyright>=1.1.398",
     "types-docker>=7.1.0.20241229",
     "pytest-timeout>=2.3.1",
+    "inline-snapshot>=0.28.0",
 ]

{dbos-1.14.0a4 → dbos-1.14.0a5}/tests/test_spans.py

@@ -1,8 +1,10 @@
-from typing import Tuple
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
 
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
+from inline_snapshot import snapshot
 from opentelemetry._logs import set_logger_provider
 from opentelemetry.sdk import trace as tracesdk
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
@@ -17,33 +19,47 @@ from dbos._tracer import dbos_tracer
 from dbos._utils import GlobalParams
 
 
+@dataclass
+class BasicSpan:
+    content: str
+    children: list["BasicSpan"] = field(default_factory=list)
+    parent_id: Optional[int] = field(repr=False, compare=False, default=None)
+
+
 def test_spans(config: DBOSConfig) -> None:
+    exporter = InMemorySpanExporter()
+    span_processor = SimpleSpanProcessor(exporter)
+    provider = tracesdk.TracerProvider()
+    provider.add_span_processor(span_processor)
+    dbos_tracer.set_provider(provider)
+
     DBOS.destroy(destroy_registry=True)
     config["otlp_attributes"] = {"foo": "bar"}
     DBOS(config=config)
     DBOS.launch()
 
+    my_tracer = provider.get_tracer("dbos")
+
     @DBOS.workflow()
     def test_workflow() -> None:
-        test_step()
-        current_span = DBOS.span
-        subspan = DBOS.tracer.start_span({"name": "a new span"}, parent=current_span)
-        # Note: DBOS.tracer.start_span() does not set the new span as the current span. So this log is still attached to the workflow span.
-        DBOS.logger.info("This is a test_workflow")
-        subspan.add_event("greeting_event", {"name": "a new event"})
-        DBOS.tracer.end_span(subspan)
+        with my_tracer.start_as_current_span(  # pyright: ignore[reportAttributeAccessIssue]
+            "manual_span"
+        ):
+            test_step()
+            current_span = DBOS.span
+            subspan = DBOS.tracer.start_span(
+                {"name": "a new span"}, parent=current_span
+            )
+            # Note: DBOS.tracer.start_span() does not set the new span as the current span. So this log is still attached to the workflow span.
+            DBOS.logger.info("This is a test_workflow")
+            subspan.add_event("greeting_event", {"name": "a new event"})
+            DBOS.tracer.end_span(subspan)
 
     @DBOS.step()
     def test_step() -> None:
         DBOS.logger.info("This is a test_step")
         return
 
-    exporter = InMemorySpanExporter()
-    span_processor = SimpleSpanProcessor(exporter)
-    provider = tracesdk.TracerProvider()
-    provider.add_span_processor(span_processor)
-    dbos_tracer.set_provider(provider)
-
     # Set up in-memory log exporter
     log_exporter = InMemoryLogExporter()  # type: ignore
     log_processor = BatchLogRecordProcessor(log_exporter)
@@ -77,9 +93,10 @@ def test_spans(config: DBOSConfig) -> None:
 
     spans = exporter.get_finished_spans()
 
-    assert len(spans) == 3
-
     for span in spans:
+        if span.name == "manual_span":
+            # Skip the manual span because it was not created by DBOS.tracer
+            continue
         assert span.attributes is not None
         assert span.attributes["applicationVersion"] == GlobalParams.app_version
         assert span.attributes["executorID"] == GlobalParams.executor_id
@@ -90,11 +107,12 @@ def test_spans(config: DBOSConfig) -> None:
 
     assert spans[0].name == test_step.__qualname__
     assert spans[1].name == "a new span"
-    assert spans[2].name == test_workflow.__qualname__
+    assert spans[3].name == test_workflow.__qualname__
 
     assert spans[0].parent.span_id == spans[2].context.span_id  # type: ignore
     assert spans[1].parent.span_id == spans[2].context.span_id  # type: ignore
-    assert spans[2].parent == None
+    assert spans[2].parent.span_id == spans[3].context.span_id  # type: ignore
+    assert spans[3].parent == None
 
     # Span ID and trace ID should match the log record
     # For pyright
@@ -105,31 +123,70 @@ def test_spans(config: DBOSConfig) -> None:
     assert logs[1].log_record.span_id == spans[2].context.span_id
     assert logs[1].log_record.trace_id == spans[2].context.trace_id
 
+    # Test the span tree structure
+    basic_spans = {
+        span.context.span_id: BasicSpan(  # pyright: ignore[reportOptionalMemberAccess]
+            content=span.name, parent_id=span.parent.span_id if span.parent else None
+        )
+        for span in spans
+    }
+    root_span = None
+    for basic_span in basic_spans.values():
+        if basic_span.parent_id is None:
+            root_span = basic_span
+        else:
+            parent_id = basic_span.parent_id
+            parent_span = basic_spans[parent_id]
+            parent_span.children.append(basic_span)
+
+    assert len(spans) == 4
+    # Make sure the span tree structure is correct
+    assert root_span == snapshot(
+        BasicSpan(
+            content="test_spans.<locals>.test_workflow",
+            children=[
+                BasicSpan(
+                    content="manual_span",
+                    children=[
+                        BasicSpan(content="test_spans.<locals>.test_step"),
+                        BasicSpan(content="a new span"),
+                    ],
+                )
+            ],
+        )
+    )
+
 
 @pytest.mark.asyncio
 async def test_spans_async(dbos: DBOS) -> None:
+    exporter = InMemorySpanExporter()
+    span_processor = SimpleSpanProcessor(exporter)
+    provider = tracesdk.TracerProvider()
+    provider.add_span_processor(span_processor)
+    dbos_tracer.set_provider(provider)
+
+    my_tracer = provider.get_tracer("dbos")
 
     @DBOS.workflow()
     async def test_workflow() -> None:
-        await test_step()
-        current_span = DBOS.span
-        subspan = DBOS.tracer.start_span({"name": "a new span"}, parent=current_span)
-        # Note: DBOS.tracer.start_span() does not set the new span as the current span. So this log is still attached to the workflow span.
-        DBOS.logger.info("This is a test_workflow")
-        subspan.add_event("greeting_event", {"name": "a new event"})
-        DBOS.tracer.end_span(subspan)
+        with my_tracer.start_as_current_span(  # pyright: ignore[reportAttributeAccessIssue]
+            "manual_span"
+        ):
+            await test_step()
+            current_span = DBOS.span
+            subspan = DBOS.tracer.start_span(
+                {"name": "a new span"}, parent=current_span
+            )
+            # Note: DBOS.tracer.start_span() does not set the new span as the current span. So this log is still attached to the workflow span.
+            DBOS.logger.info("This is a test_workflow")
+            subspan.add_event("greeting_event", {"name": "a new event"})
+            DBOS.tracer.end_span(subspan)
 
     @DBOS.step()
     async def test_step() -> None:
         DBOS.logger.info("This is a test_step")
         return
 
-    exporter = InMemorySpanExporter()
-    span_processor = SimpleSpanProcessor(exporter)
-    provider = tracesdk.TracerProvider()
-    provider.add_span_processor(span_processor)
-    dbos_tracer.set_provider(provider)
-
     # Set up in-memory log exporter
     log_exporter = InMemoryLogExporter()  # type: ignore
     log_processor = BatchLogRecordProcessor(log_exporter)
@@ -162,9 +219,12 @@ async def test_spans_async(dbos: DBOS) -> None:
 
     spans = exporter.get_finished_spans()
 
-    assert len(spans) == 3
+    assert len(spans) == 4
 
     for span in spans:
+        if span.name == "manual_span":
+            # Skip the manual span because it was not created by DBOS.tracer
+            continue
         assert span.attributes is not None
         assert span.attributes["applicationVersion"] == GlobalParams.app_version
         assert span.attributes["executorID"] == GlobalParams.executor_id
@@ -174,11 +234,12 @@ async def test_spans_async(dbos: DBOS) -> None:
 
     assert spans[0].name == test_step.__qualname__
     assert spans[1].name == "a new span"
-    assert spans[2].name == test_workflow.__qualname__
+    assert spans[3].name == test_workflow.__qualname__
 
     assert spans[0].parent.span_id == spans[2].context.span_id  # type: ignore
     assert spans[1].parent.span_id == spans[2].context.span_id  # type: ignore
-    assert spans[2].parent == None
+    assert spans[2].parent.span_id == spans[3].context.span_id  # type: ignore
+    assert spans[3].parent == None
 
     # Span ID and trace ID should match the log record
     assert spans[0].context is not None
@@ -188,6 +249,39 @@ async def test_spans_async(dbos: DBOS) -> None:
     assert logs[1].log_record.span_id == spans[2].context.span_id
     assert logs[1].log_record.trace_id == spans[2].context.trace_id
 
+    # Test the span tree structure
+    basic_spans = {
+        span.context.span_id: BasicSpan(  # pyright: ignore[reportOptionalMemberAccess]
+            content=span.name, parent_id=span.parent.span_id if span.parent else None
+        )
+        for span in spans
+    }
+    root_span = None
+    for basic_span in basic_spans.values():
+        if basic_span.parent_id is None:
+            root_span = basic_span
+        else:
+            parent_id = basic_span.parent_id
+            parent_span = basic_spans[parent_id]
+            parent_span.children.append(basic_span)
+
+    assert len(spans) == 4
+    # Make sure the span tree structure is correct
+    assert root_span == snapshot(
+        BasicSpan(
+            content="test_spans_async.<locals>.test_workflow",
+            children=[
+                BasicSpan(
+                    content="manual_span",
+                    children=[
+                        BasicSpan(content="test_spans_async.<locals>.test_step"),
+                        BasicSpan(content="a new span"),
+                    ],
+                )
+            ],
+        )
+    )
+
 
 def test_wf_fastapi(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
     dbos, app = dbos_fastapi