dbos 1.3.0a9__tar.gz → 1.4.0a1__tar.gz

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

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

{dbos-1.3.0a9 → 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.0a9 → 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.0a9 → dbos-1.4.0a1}/pyproject.toml

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

{dbos-1.3.0a9 → 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.0a9 → 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