dbos 1.11.0a3__tar.gz → 1.11.0a5__tar.gz

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

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

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

@@ -157,7 +157,7 @@ class WorkflowHandlePolling(Generic[R]):
 
 class WorkflowHandleAsyncTask(Generic[R]):
 
-    def __init__(self, workflow_id: str, task: asyncio.Task[R], dbos: "DBOS"):
+    def __init__(self, workflow_id: str, task: asyncio.Future[R], dbos: "DBOS"):
         self.workflow_id = workflow_id
         self.task = task
         self.dbos = dbos
@@ -698,7 +698,8 @@ async def start_workflow_async(
         return WorkflowHandleAsyncPolling(new_wf_id, dbos)
 
     coro = _execute_workflow_async(dbos, status, func, new_wf_ctx, *args, **kwargs)
-    task = asyncio.create_task(coro)
+    # Shield the workflow task from cancellation
+    task = asyncio.shield(asyncio.create_task(coro))
     return WorkflowHandleAsyncTask(new_wf_id, task, dbos)
 
 

{dbos-1.11.0a3 → dbos-1.11.0a5}/dbos/cli/migration.py

@@ -46,7 +46,9 @@ def grant_dbos_schema_permissions(database_url: str, role_name: str) -> None:
     )
     engine = None
     try:
-        engine = sa.create_engine(database_url)
+        engine = sa.create_engine(
+            sa.make_url(database_url).set(drivername="postgresql+psycopg")
+        )
         with engine.connect() as connection:
             connection.execution_options(isolation_level="AUTOCOMMIT")
 

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

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

{dbos-1.11.0a3 → dbos-1.11.0a5}/tests/test_async.py

@@ -595,3 +595,37 @@ async def test_main_loop(dbos: DBOS, config: DBOSConfig) -> None:
     # 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())
+
+
+@pytest.mark.asyncio
+async def test_workflow_with_task_cancellation(dbos: DBOS) -> None:
+    @DBOS.workflow()
+    async def test_workflow(duration: float) -> str:
+        await DBOS.sleep_async(duration)
+        return "completed"
+
+    # Run the workflow in an asyncio task
+    wfid = str(uuid.uuid4())
+    event = asyncio.Event()
+
+    async def run_workflow_task() -> str:
+        with SetWorkflowID(wfid):
+            handle = await DBOS.start_workflow_async(test_workflow, 1.0)
+        event.set()
+        return await handle.get_result()
+
+    task = asyncio.create_task(run_workflow_task())
+
+    # Wait for the workflow to start
+    await event.wait()
+
+    # Cancel the task
+    task.cancel()
+
+    # Verify the task was cancelled
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # Verify the workflow completes despite the task cancellation
+    handle: WorkflowHandleAsync[str] = await DBOS.retrieve_workflow_async(wfid)
+    assert await handle.get_result() == "completed"

{dbos-1.11.0a3 → dbos-1.11.0a5}/tests/test_migrate.py

@@ -11,9 +11,11 @@ def test_migrate(postgres_db_engine: sa.Engine) -> None:
     role_name = "migrate-test-role"
     role_password = "migrate_test_password"
 
-    db_url = postgres_db_engine.url.set(database=database_name).render_as_string(
-        hide_password=False
+    # Verify migration is agnostic to driver name (under the hood it uses postgresql+psycopg)
+    db_url = postgres_db_engine.url.set(database=database_name).set(
+        drivername="postgresql"
     )
+    db_url_string = db_url.render_as_string(hide_password=False)
 
     # Drop the DBOS database if it exists. Create a test role with no permissions.
     with postgres_db_engine.connect() as connection:
@@ -31,7 +33,7 @@ def test_migrate(postgres_db_engine: sa.Engine) -> None:
     # Using the admin role, create the DBOS database and verify it exists.
     # Set permissions for the test role.
     subprocess.check_call(
-        ["dbos", "migrate", "-D", db_url, "-s", db_url, "-r", role_name]
+        ["dbos", "migrate", "-D", db_url_string, "-s", db_url_string, "-r", role_name]
     )
     with postgres_db_engine.connect() as c:
         c.execution_options(isolation_level="AUTOCOMMIT")
@@ -44,9 +46,7 @@ def test_migrate(postgres_db_engine: sa.Engine) -> None:
 
     # Initialize DBOS with the test role. Verify various operations work.
     test_db_url = (
-        postgres_db_engine.url.set(database=database_name)
-        .set(username=role_name)
-        .set(password=role_password)
+        db_url.set(username=role_name).set(password=role_password)
     ).render_as_string(hide_password=False)
     DBOS.destroy(destroy_registry=True)
     config: DBOSConfig = {