dbos 1.3.0a8__tar.gz → 1.4.0a1__tar.gz

{dbos-1.3.0a8 → dbos-1.4.0a1}/PKG-INFO

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

{dbos-1.3.0a8 → dbos-1.4.0a1}/dbos/_dbos.py

@@ -100,7 +100,13 @@ from ._error import (
     DBOSNonExistentWorkflowError,
 )
 from ._event_loop import BackgroundEventLoop
-from ._logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
+from ._logger import (
+    add_otlp_to_all_loggers,
+    add_transformer_to_all_loggers,
+    config_logger,
+    dbos_logger,
+    init_logger,
+)
 from ._workflow_commands import get_workflow, list_workflow_steps
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F
@@ -514,6 +520,7 @@ class DBOS:
             for handler in dbos_logger.handlers:
                 handler.flush()
             add_otlp_to_all_loggers()
+            add_transformer_to_all_loggers()
         except Exception:
             dbos_logger.error(f"DBOS failed to launch: {traceback.format_exc()}")
             raise

{dbos-1.3.0a8 → dbos-1.4.0a1}/dbos/_event_loop.py

@@ -7,9 +7,13 @@ class BackgroundEventLoop:
     """
     This is the event loop to which DBOS submits any coroutines that are not started from within an event loop.
     In particular, coroutines submitted to queues (such as from scheduled workflows) run on this event loop.
+
+    If a main event loop is known (whether because an event loop existed in the thread that called DBOS.launch
+    or because a FastAPI event loop was detected) then coroutines are submitted there instead.
     """
 
     def __init__(self) -> None:
+        self._main_loop: Optional[asyncio.AbstractEventLoop] = None
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._thread: Optional[threading.Thread] = None
         self._running = False
@@ -19,6 +23,7 @@ class BackgroundEventLoop:
         if self._running:
             return
 
+        self.set_main_loop()
         self._thread = threading.Thread(target=self._run_event_loop, daemon=True)
         self._thread.start()
         self._ready.wait()  # Wait until the loop is running
@@ -58,10 +63,23 @@ class BackgroundEventLoop:
         await asyncio.gather(*tasks, return_exceptions=True)
         self._loop.stop()
 
+    def set_main_loop(self) -> None:
+        """
+        Set the main loop to the currently running event loop.
+        Should be called from the main thread.
+        """
+        try:
+            self._main_loop = asyncio.get_running_loop()
+        except:
+            # There's no running event loop to set
+            pass
+
     T = TypeVar("T")
 
     def submit_coroutine(self, coro: Coroutine[Any, Any, T]) -> T:
         """Submit a coroutine to the background event loop"""
+        if self._main_loop is not None and self._main_loop.is_running():
+            return asyncio.run_coroutine_threadsafe(coro, self._main_loop).result()
         if self._loop is None:
             raise RuntimeError("Event loop not started")
         return asyncio.run_coroutine_threadsafe(coro, self._loop).result()

{dbos-1.3.0a8 → dbos-1.4.0a1}/dbos/_fastapi.py

@@ -44,11 +44,10 @@ class LifespanMiddleware:
         if scope["type"] == "lifespan":
 
             async def wrapped_send(message: MutableMapping[str, Any]) -> None:
-                if (
-                    message["type"] == "lifespan.startup.complete"
-                    and not self.dbos._launched
-                ):
-                    self.dbos._launch()
+                if message["type"] == "lifespan.startup.complete":
+                    self.dbos._background_event_loop.set_main_loop()
+                    if not self.dbos._launched:
+                        self.dbos._launch()
                 elif message["type"] == "lifespan.shutdown.complete":
                     self.dbos._destroy()
                 await send(message)

{dbos-1.3.0a8 → dbos-1.4.0a1}/dbos/_logger.py

@@ -16,7 +16,7 @@ if TYPE_CHECKING:
     from ._dbos_config import ConfigFile
 
 dbos_logger = logging.getLogger("dbos")
-_otlp_handler, _otlp_transformer = None, None
+_otlp_handler, _dbos_log_transformer = None, None
 
 
 class DBOSLogTransformer(logging.Filter):
@@ -93,25 +93,31 @@ def config_logger(config: "ConfigFile") -> None:
         global _otlp_handler
         _otlp_handler = LoggingHandler(logger_provider=log_provider)
 
-        # Attach DBOS-specific attributes to all log entries.
-        global _otlp_transformer
-        _otlp_transformer = DBOSLogTransformer()
-
         # Direct DBOS logs to OTLP
         dbos_logger.addHandler(_otlp_handler)
-        dbos_logger.addFilter(_otlp_transformer)
+
+    # Attach DBOS-specific attributes to all log entries.
+    global _dbos_log_transformer
+    _dbos_log_transformer = DBOSLogTransformer()
+    dbos_logger.addFilter(_dbos_log_transformer)
 
 
 def add_otlp_to_all_loggers() -> None:
-    if _otlp_handler is not None and _otlp_transformer is not None:
+    if _otlp_handler is not None:
         root = logging.root
-
         root.addHandler(_otlp_handler)
-        root.addFilter(_otlp_transformer)
-
         for logger_name in root.manager.loggerDict:
             if logger_name != dbos_logger.name:
                 logger = logging.getLogger(logger_name)
                 if not logger.propagate:
                     logger.addHandler(_otlp_handler)
-                logger.addFilter(_otlp_transformer)
+
+
+def add_transformer_to_all_loggers() -> None:
+    if _dbos_log_transformer is not None:
+        root = logging.root
+        root.addFilter(_dbos_log_transformer)
+        for logger_name in root.manager.loggerDict:
+            if logger_name != dbos_logger.name:
+                logger = logging.getLogger(logger_name)
+                logger.addFilter(_dbos_log_transformer)

{dbos-1.3.0a8 → dbos-1.4.0a1}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "1.3.0a8"
+version = "1.4.0a1"
 
 [project.license]
 text = "MIT"

{dbos-1.3.0a8 → dbos-1.4.0a1}/tests/test_async.py

@@ -7,7 +7,14 @@ import pytest
 import sqlalchemy as sa
 
 # Public API
-from dbos import DBOS, Queue, SetWorkflowID, SetWorkflowTimeout, WorkflowHandleAsync
+from dbos import (
+    DBOS,
+    DBOSConfig,
+    Queue,
+    SetWorkflowID,
+    SetWorkflowTimeout,
+    WorkflowHandleAsync,
+)
 from dbos._context import assert_current_dbos_context
 from dbos._dbos import WorkflowHandle
 from dbos._dbos_config import ConfigFile
@@ -564,3 +571,21 @@ async def test_max_parallel_workflows(dbos: DBOS) -> None:
     assert (
         end_time - begin_time < 10
     ), "All enqueued tasks should complete in less than 10 seconds"
+
+
+@pytest.mark.asyncio
+async def test_main_loop(dbos: DBOS, config: DBOSConfig) -> None:
+    DBOS.destroy(destroy_registry=True)
+    dbos = DBOS(config=config)
+    DBOS.launch()
+
+    queue = Queue("queue")
+
+    @DBOS.workflow()
+    async def test_workflow() -> int:
+        await DBOS.sleep_async(0.1)
+        return id(asyncio.get_running_loop())
+
+    # Verify the enqueued task is submitted into the main event loop
+    handle = await queue.enqueue_async(test_workflow)
+    assert await handle.get_result() == id(asyncio.get_running_loop())

{dbos-1.3.0a8 → dbos-1.4.0a1}/tests/test_fastapi.py

@@ -12,7 +12,7 @@ from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
 # Public API
-from dbos import DBOS, DBOSConfig
+from dbos import DBOS, DBOSConfig, Queue
 
 # Private API because this is a unit test
 from dbos._context import assert_current_dbos_context
@@ -155,10 +155,21 @@ async def test_custom_lifespan(
     DBOS.destroy()
     DBOS(fastapi=app, config=config)
 
+    queue = Queue("queue")
+
     @app.get("/")
     @DBOS.workflow()
     async def resource_workflow() -> Any:
-        return {"resource": resource}
+        handle = await queue.enqueue_async(queue_workflow)
+        return {
+            "resource": resource,
+            "loop": id(asyncio.get_event_loop()),
+            "queue_loop": await handle.get_result(),
+        }
+
+    @DBOS.workflow()
+    async def queue_workflow() -> int:
+        return id(asyncio.get_event_loop())
 
     uvicorn_config = uvicorn.Config(
         app=app, host="127.0.0.1", port=port, log_level="error"
@@ -172,6 +183,9 @@ async def test_custom_lifespan(
     async with httpx.AsyncClient() as client:
         r = await client.get(f"http://127.0.0.1:{port}")
         assert r.json()["resource"] == 1
+        # Verify that both the FastAPI and enqueued workflows run in the main event loop
+        assert r.json()["loop"] == id(asyncio.get_event_loop())
+        assert r.json()["queue_loop"] == id(asyncio.get_event_loop())
 
     server.should_exit = True
     await server_task