Temp Workflows Shouldn't Create Spans (#169)

Author: kraftp

dbos/_context.py

@@ -98,18 +98,27 @@ def assign_workflow_id(self) -> str:
             wfid = str(uuid.uuid4())
         return wfid
 
-    def start_workflow(self, wfid: Optional[str], attributes: TracedAttributes) -> None:
+    def start_workflow(
+        self,
+        wfid: Optional[str],
+        attributes: TracedAttributes,
+        is_temp_workflow: bool = False,
+    ) -> None:
         if wfid is None or len(wfid) == 0:
             wfid = self.assign_workflow_id()
             self.id_assigned_for_next_workflow = ""
         self.workflow_id = wfid
         self.function_id = 0
-        self._start_span(attributes)
+        if not is_temp_workflow:
+            self._start_span(attributes)
 
-    def end_workflow(self, exc_value: Optional[BaseException]) -> None:
+    def end_workflow(
+        self, exc_value: Optional[BaseException], is_temp_workflow: bool = False
+    ) -> None:
         self.workflow_id = ""
         self.function_id = -1
-        self._end_span(exc_value)
+        if not is_temp_workflow:
+            self._end_span(exc_value)
 
     def is_within_workflow(self) -> bool:
         return len(self.workflow_id) > 0
@@ -349,6 +358,7 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
     def __init__(self, attributes: TracedAttributes) -> None:
         self.created_ctx = False
         self.attributes = attributes
+        self.is_temp_workflow = attributes["name"] == "temp_wf"
 
     def __enter__(self) -> DBOSContext:
         # Code to create a basic context
@@ -359,7 +369,7 @@ def __enter__(self) -> DBOSContext:
             _set_local_dbos_context(ctx)
         assert not ctx.is_within_workflow()
         ctx.start_workflow(
-            None, self.attributes
+            None, self.attributes, self.is_temp_workflow
         )  # Will get from the context's next workflow ID
         return ctx
 
@@ -371,7 +381,7 @@ def __exit__(
     ) -> Literal[False]:
         ctx = assert_current_dbos_context()
         assert ctx.is_within_workflow()
-        ctx.end_workflow(exc_value)
+        ctx.end_workflow(exc_value, self.is_temp_workflow)
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

tests/test_spans.py

@@ -0,0 +1,80 @@
+from typing import Tuple
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from opentelemetry.sdk import trace as tracesdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from dbos import DBOS
+from dbos._tracer import dbos_tracer
+
+
+def test_spans(dbos: DBOS) -> None:
+
+    @DBOS.workflow()
+    def test_workflow() -> None:
+        test_step()
+
+    @DBOS.step()
+    def test_step() -> None:
+        return
+
+    exporter = InMemorySpanExporter()
+    span_processor = SimpleSpanProcessor(exporter)
+    provider = tracesdk.TracerProvider()
+    provider.add_span_processor(span_processor)
+    dbos_tracer.set_provider(provider)
+
+    test_workflow()
+    test_step()
+
+    spans = exporter.get_finished_spans()
+
+    assert len(spans) == 3
+
+    for span in spans:
+        assert span.attributes is not None
+        assert span.context is not None
+
+    assert spans[0].name == test_step.__name__
+    assert spans[1].name == test_workflow.__name__
+    assert spans[2].name == test_step.__name__
+
+    assert spans[0].parent.span_id == spans[1].context.span_id  # type: ignore
+    assert spans[1].parent == None
+    assert spans[2].parent == None
+
+
+def test_temp_wf_fastapi(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
+    _, app = dbos_fastapi
+
+    @app.get("/step")
+    @DBOS.step()
+    def test_step_endpoint() -> str:
+        return "test"
+
+    exporter = InMemorySpanExporter()
+    span_processor = SimpleSpanProcessor(exporter)
+    provider = tracesdk.TracerProvider()
+    provider.add_span_processor(span_processor)
+    dbos_tracer.set_provider(provider)
+
+    client = TestClient(app)
+    response = client.get("/step")
+    assert response.status_code == 200
+    assert response.text == '"test"'
+
+    spans = exporter.get_finished_spans()
+
+    assert len(spans) == 2
+
+    for span in spans:
+        assert span.attributes is not None
+        assert span.context is not None
+
+    assert spans[0].name == test_step_endpoint.__name__
+    assert spans[1].name == "/step"
+
+    assert spans[0].parent.span_id == spans[1].context.span_id  # type:ignore
+    assert spans[1].parent == None