dbos 1.14.0a3__tar.gz → 1.14.0a5__tar.gz

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

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

{dbos-1.14.0a3 → 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.0a3 → dbos-1.14.0a5}/dbos/_core.py

@@ -56,6 +56,7 @@ from ._error import (
     DBOSWorkflowConflictIDError,
     DBOSWorkflowFunctionNotFoundError,
 )
+from ._logger import dbos_logger
 from ._registrations import (
     DEFAULT_MAX_RECOVERY_ATTEMPTS,
     get_config_name,
@@ -96,6 +97,14 @@ F = TypeVar("F", bound=Callable[..., Any])
 TEMP_SEND_WF_NAME = "<temp>.temp_send_workflow"
 
 
+def check_is_in_coroutine() -> bool:
+    try:
+        asyncio.get_running_loop()
+        return True
+    except RuntimeError:
+        return False
+
+
 class WorkflowHandleFuture(Generic[R]):
 
     def __init__(self, workflow_id: str, future: Future[R], dbos: "DBOS"):
@@ -828,6 +837,11 @@ def workflow_wrapper(
             dbos._sys_db.record_get_result(workflow_id, serialized_r, None)
             return r
 
+        if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
+            dbos_logger.warning(
+                f"Sync workflow ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
+            )
+
         outcome = (
             wfOutcome.wrap(init_wf, dbos=dbos)
             .also(DBOSAssumeRole(rr))
@@ -957,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",
@@ -1009,6 +1023,10 @@ def decorate_transaction(
                 assert (
                     ctx.is_workflow()
                 ), "Transactions must be called from within workflows"
+                if check_is_in_coroutine():
+                    dbos_logger.warning(
+                        f"Transaction function ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Use asyncio.to_thread instead."
+                    )
                 with DBOSAssumeRole(rr):
                     return invoke_tx(*args, **kwargs)
             else:
@@ -1072,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",
@@ -1153,6 +1171,10 @@ def decorate_step(
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
+                dbos_logger.warning(
+                    f"Sync step ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
+                )
             # If the step is called from a workflow, run it as a step.
             # Otherwise, run it as a normal function.
             ctx = get_local_dbos_context()

{dbos-1.14.0a3 → 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.0a3 → 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.0a3 → 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.0a3 → dbos-1.14.0a5}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "1.14.0a3"
+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.0a3 → dbos-1.14.0a5}/tests/test_async.py

@@ -5,6 +5,9 @@ from typing import List, Optional
 
 import pytest
 import sqlalchemy as sa
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, InMemoryLogExporter
 
 # Public API
 from dbos import (
@@ -19,6 +22,8 @@ from dbos._context import assert_current_dbos_context
 from dbos._dbos import WorkflowHandle
 from dbos._dbos_config import ConfigFile
 from dbos._error import DBOSAwaitedWorkflowCancelledError, DBOSException
+from dbos._logger import dbos_logger
+from dbos._registrations import get_dbos_func_name
 
 
 @pytest.mark.asyncio
@@ -31,7 +36,7 @@ async def test_async_workflow(dbos: DBOS) -> None:
     async def test_workflow(var1: str, var2: str) -> str:
         nonlocal wf_counter
         wf_counter += 1
-        res1 = test_transaction(var1)
+        res1 = await asyncio.to_thread(test_transaction, var1)
         res2 = await test_step(var2)
         DBOS.logger.info("I'm test_workflow")
         return res1 + res2
@@ -88,7 +93,7 @@ async def test_async_step(dbos: DBOS) -> None:
     async def test_workflow(var1: str, var2: str) -> str:
         nonlocal wf_counter
         wf_counter += 1
-        res1 = test_transaction(var1)
+        res1 = await asyncio.to_thread(test_transaction, var1)
         res2 = await test_step(var2)
         DBOS.logger.info("I'm test_workflow")
         return res1 + res2
@@ -325,6 +330,7 @@ def test_async_tx_raises(config: ConfigFile) -> None:
         async def test_async_tx() -> None:
             pass
 
+    assert "is a coroutine function" in str(exc_info.value)
     # destroy call needed to avoid "functions were registered but DBOS() was not called" warning
     DBOS.destroy(destroy_registry=True)
 
@@ -343,12 +349,12 @@ async def test_start_workflow_async(dbos: DBOS) -> None:
         wf_el_id = id(asyncio.get_running_loop())
         nonlocal wf_counter
         wf_counter += 1
-        res2 = test_step(var2)
+        res2 = await test_step(var2)
         DBOS.logger.info("I'm test_workflow")
         return var1 + res2
 
     @DBOS.step()
-    def test_step(var: str) -> str:
+    async def test_step(var: str) -> str:
         nonlocal step_el_id
         step_el_id = id(asyncio.get_running_loop())
         nonlocal step_counter
@@ -605,3 +611,83 @@ async def test_workflow_with_task_cancellation(dbos: DBOS) -> None:
     # Verify the workflow completes despite the task cancellation
     handle: WorkflowHandleAsync[str] = await DBOS.retrieve_workflow_async(wfid)
     assert await handle.get_result() == "completed"
+
+
+@pytest.mark.asyncio
+async def test_check_async_violation(dbos: DBOS) -> None:
+    # Set up in-memory log exporter
+    log_exporter = InMemoryLogExporter()  # type: ignore
+    log_processor = BatchLogRecordProcessor(log_exporter)
+    log_provider = LoggerProvider()
+    log_provider.add_log_record_processor(log_processor)
+    set_logger_provider(log_provider)
+    dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
+
+    @DBOS.workflow()
+    def sync_workflow() -> str:
+        return "sync"
+
+    @DBOS.step()
+    def sync_step() -> str:
+        return "step"
+
+    @DBOS.workflow()
+    async def async_workflow_sync_step() -> str:
+        return sync_step()
+
+    @DBOS.transaction()
+    def sync_transaction() -> str:
+        return "txn"
+
+    @DBOS.workflow()
+    async def async_workflow_sync_txn() -> str:
+        return sync_transaction()
+
+    # Call a sync workflow should log a warning
+    sync_workflow()
+
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Sync workflow ({get_dbos_func_name(sync_workflow)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()
+
+    # Call a sync step from within an async workflow should log a warning
+    await async_workflow_sync_step()
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Sync step ({get_dbos_func_name(sync_step)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()
+
+    # Directly call a sync step should log a warning
+    sync_step()
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Sync step ({get_dbos_func_name(sync_step)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()
+
+    # Call a sync transaction from within an async workflow should log a warning
+    await async_workflow_sync_txn()
+    log_processor.force_flush(timeout_millis=5000)
+    logs = log_exporter.get_finished_logs()
+    assert len(logs) == 1
+    assert (
+        logs[0].log_record.body is not None
+        and f"Transaction function ({get_dbos_func_name(sync_transaction)}) shouldn't be invoked from within another async function."
+        in logs[0].log_record.body
+    )
+    log_exporter.clear()

{dbos-1.14.0a3 → dbos-1.14.0a5}/tests/test_async_workflow_management.py

@@ -148,7 +148,7 @@ async def test_fork_workflow_async(dbos: DBOS) -> None:
 
     wfid = str(uuid.uuid4())
     with SetWorkflowID(wfid):
-        assert simple_workflow(input_val) == output
+        assert await asyncio.to_thread(simple_workflow, input_val) == output
 
     assert step_one_count == 1
     assert step_two_count == 1

{dbos-1.14.0a3 → dbos-1.14.0a5}/tests/test_dbos.py

@@ -1,5 +1,6 @@
 # mypy: disable-error-code="no-redef"
 
+import asyncio
 import datetime
 import logging
 import os
@@ -1668,17 +1669,17 @@ async def test_step_without_dbos(dbos: DBOS, config: DBOSConfig) -> None:
         assert DBOS.workflow_id is None
         return x
 
-    assert step(5) == 5
+    assert await asyncio.to_thread(step, 5) == 5
     assert await async_step(5) == 5
 
     DBOS(config=config)
 
-    assert step(5) == 5
+    assert await asyncio.to_thread(step, 5) == 5
     assert await async_step(5) == 5
 
     DBOS.launch()
 
-    assert step(5) == 5
+    assert await asyncio.to_thread(step, 5) == 5
     assert await async_step(5) == 5
 
     assert len(DBOS.list_workflows()) == 0

{dbos-1.14.0a3 → 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

{dbos-1.14.0a3 → dbos-1.14.0a5}/tests/test_workflow_introspection.py

@@ -828,18 +828,18 @@ async def test_callchild_first_asyncio(dbos: DBOS) -> None:
     async def parentWorkflow() -> str:
         handle = await dbos.start_workflow_async(child_workflow)
         child_id = await handle.get_result()
-        stepOne()
-        stepTwo()
+        await stepOne()
+        await stepTwo()
         return child_id
 
     @DBOS.step()
-    def stepOne() -> str:
+    async def stepOne() -> str:
         workflow_id = DBOS.workflow_id
         assert workflow_id is not None
         return workflow_id
 
     @DBOS.step()
-    def stepTwo() -> None:
+    async def stepTwo() -> None:
         return
 
     @DBOS.workflow()