dbos 0.20.0a3__tar.gz → 0.20.0a5__tar.gz

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

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

{dbos-0.20.0a3 → dbos-0.20.0a5}/dbos/_core.py

@@ -488,6 +488,22 @@ def start_workflow(
     return WorkflowHandleFuture(new_wf_id, future, dbos)
 
 
+if sys.version_info < (3, 12):
+
+    def _mark_coroutine(func: Callable[P, R]) -> Callable[P, R]:
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> R:
+            return await func(*args, **kwargs)  # type: ignore
+
+        return async_wrapper  # type: ignore
+
+else:
+
+    def _mark_coroutine(func: Callable[P, R]) -> Callable[P, R]:
+        inspect.markcoroutinefunction(func)
+        return func
+
+
 def workflow_wrapper(
     dbosreg: "DBOSRegistry",
     func: Callable[P, R],
@@ -548,7 +564,7 @@ def workflow_wrapper(
         )
         return outcome()  # type: ignore
 
-    return wrapper
+    return _mark_coroutine(wrapper) if inspect.iscoroutinefunction(func) else wrapper
 
 
 def decorate_workflow(
@@ -838,6 +854,10 @@ def decorate_step(
                 assert tempwf
                 return tempwf(*args, **kwargs)
 
+        wrapper = (
+            _mark_coroutine(wrapper) if inspect.iscoroutinefunction(func) else wrapper  # type: ignore
+        )
+
         def temp_wf_sync(*args: Any, **kwargs: Any) -> Any:
             return wrapper(*args, **kwargs)
 

{dbos-0.20.0a3 → dbos-0.20.0a5}/dbos/_fastapi.py

@@ -1,10 +1,10 @@
 import uuid
-from typing import Any, Callable, cast
+from typing import Any, Callable, MutableMapping, cast
 
 from fastapi import FastAPI
 from fastapi import Request as FastAPIRequest
 from fastapi.responses import JSONResponse
-from starlette.types import ASGIApp, Message, Receive, Scope, Send
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 from . import DBOS
 from ._context import (
@@ -61,15 +61,16 @@ class LifespanMiddleware:
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] == "lifespan":
-            while True:
-                message = await receive()
-                if message["type"] == "lifespan.startup":
+
+            async def wrapped_send(message: MutableMapping[str, Any]) -> None:
+                if message["type"] == "lifespan.startup.complete":
                     self.dbos._launch()
-                    await send({"type": "lifespan.startup.complete"})
-                elif message["type"] == "lifespan.shutdown":
+                elif message["type"] == "lifespan.shutdown.complete":
                     self.dbos._destroy()
-                    await send({"type": "lifespan.shutdown.complete"})
-                    break
+                await send(message)
+
+            # Call the original app with our wrapped functions
+            await self.app(scope, receive, wrapped_send)
         else:
             await self.app(scope, receive, send)
 

{dbos-0.20.0a3 → dbos-0.20.0a5}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.20.0a3"
+version = "0.20.0a5"
 
 [project.license]
 text = "MIT"

{dbos-0.20.0a3 → dbos-0.20.0a5}/tests/conftest.py

@@ -118,15 +118,7 @@ def dbos_fastapi(
 ) -> Generator[Tuple[DBOS, FastAPI], Any, None]:
     DBOS.destroy()
     app = FastAPI()
-
-    # ignore the on_event deprecation warnings
-    with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore",
-            category=DeprecationWarning,
-            message=r"\s*on_event is deprecated, use lifespan event handlers instead\.",
-        )
-        dbos = DBOS(fastapi=app, config=config)
+    dbos = DBOS(fastapi=app, config=config)
 
     # This is for test convenience.
     #    Usually fastapi itself does launch, but we are not completing the fastapi lifecycle

{dbos-0.20.0a3 → dbos-0.20.0a5}/tests/test_fastapi.py

@@ -1,14 +1,18 @@
+import asyncio
 import logging
 import uuid
-from typing import Tuple
+from contextlib import asynccontextmanager
+from typing import Any, Tuple
 
+import httpx
 import pytest
 import sqlalchemy as sa
+import uvicorn
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
 # Public API
-from dbos import DBOS
+from dbos import DBOS, ConfigFile
 
 # Private API because this is a unit test
 from dbos._context import assert_current_dbos_context
@@ -157,3 +161,73 @@ def test_endpoint_recovery(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
     workflow_handles = DBOS.recover_pending_workflows()
     assert len(workflow_handles) == 1
     assert workflow_handles[0].get_result() == ("a", wfuuid)
+
+
+@pytest.mark.asyncio
+async def test_custom_lifespan(
+    config: ConfigFile, cleanup_test_databases: None
+) -> None:
+    resource = None
+    port = 8000
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI) -> Any:
+        nonlocal resource
+        resource = 1
+        yield
+        resource = None
+
+    app = FastAPI(lifespan=lifespan)
+
+    DBOS.destroy()
+    DBOS(fastapi=app, config=config)
+
+    @app.get("/")
+    @DBOS.workflow()
+    async def resource_workflow() -> Any:
+        return {"resource": resource}
+
+    uvicorn_config = uvicorn.Config(
+        app=app, host="127.0.0.1", port=port, log_level="error"
+    )
+    server = uvicorn.Server(config=uvicorn_config)
+
+    # Run server in background task
+    server_task = asyncio.create_task(server.serve())
+    await asyncio.sleep(0.2)  # Give server time to start
+
+    async with httpx.AsyncClient() as client:
+        r = await client.get(f"http://127.0.0.1:{port}")
+        assert r.json()["resource"] == 1
+
+    server.should_exit = True
+    await server_task
+    assert resource is None
+
+
+def test_stacked_decorators_wf(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
+    dbos, app = dbos_fastapi
+    client = TestClient(app)
+
+    @app.get("/endpoint/{var1}/{var2}")
+    @DBOS.workflow()
+    async def test_endpoint(var1: str, var2: str) -> str:
+        return f"{var1}, {var2}!"
+
+    response = client.get("/endpoint/plums/deify")
+    assert response.status_code == 200
+    assert response.text == '"plums, deify!"'
+
+
+def test_stacked_decorators_step(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
+    dbos, app = dbos_fastapi
+    client = TestClient(app)
+
+    @app.get("/endpoint/{var1}/{var2}")
+    @DBOS.step()
+    async def test_endpoint(var1: str, var2: str) -> str:
+        return f"{var1}, {var2}!"
+
+    response = client.get("/endpoint/plums/deify")
+    assert response.status_code == 200
+    assert response.text == '"plums, deify!"'