dbos 1.8.0a3__tar.gz → 1.8.0a5__tar.gz

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

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

{dbos-1.8.0a3 → dbos-1.8.0a5}/dbos/_client.py

@@ -13,7 +13,7 @@ else:
 
 from dbos import _serialization
 from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
-from dbos._dbos_config import is_valid_database_url
+from dbos._dbos_config import get_system_database_url, is_valid_database_url
 from dbos._error import DBOSException, DBOSNonExistentWorkflowError
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
 from dbos._serialization import WorkflowInputs
@@ -97,17 +97,28 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
 
 
 class DBOSClient:
-    def __init__(self, database_url: str, *, system_database: Optional[str] = None):
+    def __init__(
+        self,
+        database_url: str,
+        *,
+        system_database_url: Optional[str] = None,
+        system_database: Optional[str] = None,
+    ):
         assert is_valid_database_url(database_url)
         # We only create database connections but do not run migrations
         self._sys_db = SystemDatabase(
-            database_url=database_url,
+            system_database_url=get_system_database_url(
+                {
+                    "system_database_url": system_database_url,
+                    "database_url": database_url,
+                    "database": {"sys_db_name": system_database},
+                }
+            ),
             engine_kwargs={
                 "pool_timeout": 30,
                 "max_overflow": 0,
                 "pool_size": 2,
             },
-            sys_db_name=system_database,
         )
         self._sys_db.check_connection()
         self._app_db = ApplicationDatabase(

{dbos-1.8.0a3 → dbos-1.8.0a5}/dbos/_dbos.py

@@ -91,6 +91,7 @@ from ._context import (
 from ._dbos_config import (
     ConfigFile,
     DBOSConfig,
+    get_system_database_url,
     overwrite_config,
     process_config,
     translate_dbos_config_to_config_file,
@@ -424,9 +425,8 @@ class DBOS:
             assert self._config["database_url"] is not None
             assert self._config["database"]["sys_db_engine_kwargs"] is not None
             self._sys_db_field = SystemDatabase(
-                database_url=self._config["database_url"],
+                system_database_url=get_system_database_url(self._config),
                 engine_kwargs=self._config["database"]["sys_db_engine_kwargs"],
-                sys_db_name=self._config["database"]["sys_db_name"],
                 debug_mode=debug_mode,
             )
             assert self._config["database"]["db_engine_kwargs"] is not None
@@ -966,6 +966,12 @@ class DBOS:
             fn, "DBOS.cancelWorkflow"
         )
 
+    @classmethod
+    async def cancel_workflow_async(cls, workflow_id: str) -> None:
+        """Cancel a workflow by ID."""
+        await cls._configure_asyncio_thread_pool()
+        await asyncio.to_thread(cls.cancel_workflow, workflow_id)
+
     @classmethod
     async def _configure_asyncio_thread_pool(cls) -> None:
         """
@@ -987,11 +993,23 @@ class DBOS:
         _get_dbos_instance()._sys_db.call_function_as_step(fn, "DBOS.resumeWorkflow")
         return cls.retrieve_workflow(workflow_id)
 
+    @classmethod
+    async def resume_workflow_async(cls, workflow_id: str) -> WorkflowHandleAsync[Any]:
+        """Resume a workflow by ID."""
+        await cls._configure_asyncio_thread_pool()
+        await asyncio.to_thread(cls.resume_workflow, workflow_id)
+        return await cls.retrieve_workflow_async(workflow_id)
+
     @classmethod
     def restart_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Restart a workflow with a new workflow ID"""
         return cls.fork_workflow(workflow_id, 1)
 
+    @classmethod
+    async def restart_workflow_async(cls, workflow_id: str) -> WorkflowHandleAsync[Any]:
+        """Restart a workflow with a new workflow ID"""
+        return await cls.fork_workflow_async(workflow_id, 1)
+
     @classmethod
     def fork_workflow(
         cls,
@@ -1017,6 +1035,23 @@ class DBOS:
         )
         return cls.retrieve_workflow(new_id)
 
+    @classmethod
+    async def fork_workflow_async(
+        cls,
+        workflow_id: str,
+        start_step: int,
+        *,
+        application_version: Optional[str] = None,
+    ) -> WorkflowHandleAsync[Any]:
+        """Restart a workflow with a new workflow ID from a specific step"""
+        await cls._configure_asyncio_thread_pool()
+        new_id = await asyncio.to_thread(
+            lambda: cls.fork_workflow(
+                workflow_id, start_step, application_version=application_version
+            ).get_workflow_id()
+        )
+        return await cls.retrieve_workflow_async(new_id)
+
     @classmethod
     def list_workflows(
         cls,
@@ -1057,6 +1092,42 @@ class DBOS:
             fn, "DBOS.listWorkflows"
         )
 
+    @classmethod
+    async def list_workflows_async(
+        cls,
+        *,
+        workflow_ids: Optional[List[str]] = None,
+        status: Optional[Union[str, List[str]]] = None,
+        start_time: Optional[str] = None,
+        end_time: Optional[str] = None,
+        name: Optional[str] = None,
+        app_version: Optional[str] = None,
+        user: Optional[str] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        sort_desc: bool = False,
+        workflow_id_prefix: Optional[str] = None,
+        load_input: bool = True,
+        load_output: bool = True,
+    ) -> List[WorkflowStatus]:
+        await cls._configure_asyncio_thread_pool()
+        return await asyncio.to_thread(
+            cls.list_workflows,
+            workflow_ids=workflow_ids,
+            status=status,
+            start_time=start_time,
+            end_time=end_time,
+            name=name,
+            app_version=app_version,
+            user=user,
+            limit=limit,
+            offset=offset,
+            sort_desc=sort_desc,
+            workflow_id_prefix=workflow_id_prefix,
+            load_input=load_input,
+            load_output=load_output,
+        )
+
     @classmethod
     def list_queued_workflows(
         cls,
@@ -1089,6 +1160,34 @@ class DBOS:
             fn, "DBOS.listQueuedWorkflows"
         )
 
+    @classmethod
+    async def list_queued_workflows_async(
+        cls,
+        *,
+        queue_name: Optional[str] = None,
+        status: Optional[Union[str, List[str]]] = None,
+        start_time: Optional[str] = None,
+        end_time: Optional[str] = None,
+        name: Optional[str] = None,
+        limit: Optional[int] = None,
+        offset: Optional[int] = None,
+        sort_desc: bool = False,
+        load_input: bool = True,
+    ) -> List[WorkflowStatus]:
+        await cls._configure_asyncio_thread_pool()
+        return await asyncio.to_thread(
+            cls.list_queued_workflows,
+            queue_name=queue_name,
+            status=status,
+            start_time=start_time,
+            end_time=end_time,
+            name=name,
+            limit=limit,
+            offset=offset,
+            sort_desc=sort_desc,
+            load_input=load_input,
+        )
+
     @classmethod
     def list_workflow_steps(cls, workflow_id: str) -> List[StepInfo]:
         def fn() -> List[StepInfo]:
@@ -1100,6 +1199,11 @@ class DBOS:
             fn, "DBOS.listWorkflowSteps"
         )
 
+    @classmethod
+    async def list_workflow_steps_async(cls, workflow_id: str) -> List[StepInfo]:
+        await cls._configure_asyncio_thread_pool()
+        return await asyncio.to_thread(cls.list_workflow_steps, workflow_id)
+
     @classproperty
     def logger(cls) -> Logger:
         """Return the DBOS `Logger` for the current context."""

{dbos-1.8.0a3 → dbos-1.8.0a5}/dbos/_dbos_config.py

@@ -23,7 +23,8 @@ class DBOSConfig(TypedDict, total=False):
     Attributes:
         name (str): Application name
         database_url (str): Database connection string
-        sys_db_name (str): System database name
+        system_database_url (str): Connection string for the system database (if different from the application database)
+        sys_db_name (str): System database name (deprecated)
         sys_db_pool_size (int): System database pool size
         db_engine_kwargs (Dict[str, Any]): SQLAlchemy engine kwargs (See https://docs.sqlalchemy.org/en/20/core/engines.html#sqlalchemy.create_engine)
         log_level (str): Log level
@@ -36,6 +37,7 @@ class DBOSConfig(TypedDict, total=False):
 
     name: str
     database_url: Optional[str]
+    system_database_url: Optional[str]
     sys_db_name: Optional[str]
     sys_db_pool_size: Optional[int]
     db_engine_kwargs: Optional[Dict[str, Any]]
@@ -111,6 +113,7 @@ class ConfigFile(TypedDict, total=False):
     runtimeConfig: RuntimeConfig
     database: DatabaseConfig
     database_url: Optional[str]
+    system_database_url: Optional[str]
     telemetry: Optional[TelemetryConfig]
     env: Dict[str, str]
 
@@ -136,6 +139,8 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
 
     if "database_url" in config:
         translated_config["database_url"] = config.get("database_url")
+    if "system_database_url" in config:
+        translated_config["system_database_url"] = config.get("system_database_url")
 
     # Runtime config
     translated_config["runtimeConfig"] = {"run_admin_server": True}
@@ -488,6 +493,8 @@ def overwrite_config(provided_config: ConfigFile) -> ConfigFile:
             "DBOS_DATABASE_URL environment variable is not set. This is required to connect to the database."
         )
     provided_config["database_url"] = db_url
+    if "system_database_url" in provided_config:
+        del provided_config["system_database_url"]
 
     # Telemetry config
     if "telemetry" not in provided_config or provided_config["telemetry"] is None:
@@ -537,3 +544,19 @@ def overwrite_config(provided_config: ConfigFile) -> ConfigFile:
         del provided_config["env"]
 
     return provided_config
+
+
+def get_system_database_url(config: ConfigFile) -> str:
+    if "system_database_url" in config and config["system_database_url"] is not None:
+        return config["system_database_url"]
+    else:
+        assert config["database_url"] is not None
+        app_db_url = make_url(config["database_url"])
+        if config["database"].get("sys_db_name") is not None:
+            sys_db_name = config["database"]["sys_db_name"]
+        else:
+            assert app_db_url.database is not None
+            sys_db_name = app_db_url.database + SystemSchema.sysdb_suffix
+        return app_db_url.set(database=sys_db_name).render_as_string(
+            hide_password=False
+        )

{dbos-1.8.0a3 → dbos-1.8.0a5}/dbos/_sys_db.py

@@ -331,22 +331,15 @@ class SystemDatabase:
     def __init__(
         self,
         *,
-        database_url: str,
+        system_database_url: str,
         engine_kwargs: Dict[str, Any],
-        sys_db_name: Optional[str] = None,
         debug_mode: bool = False,
     ):
         # Set driver
-        system_db_url = sa.make_url(database_url).set(drivername="postgresql+psycopg")
-        # Resolve system database name
-        sysdb_name = sys_db_name
-        if not sysdb_name:
-            assert system_db_url.database is not None
-            sysdb_name = system_db_url.database + SystemSchema.sysdb_suffix
-        system_db_url = system_db_url.set(database=sysdb_name)
+        url = sa.make_url(system_database_url).set(drivername="postgresql+psycopg")
 
         self.engine = sa.create_engine(
-            system_db_url,
+            url,
             **engine_kwargs,
         )
         self._engine_kwargs = engine_kwargs

{dbos-1.8.0a3 → dbos-1.8.0a5}/dbos/cli/cli.py

@@ -21,6 +21,7 @@ from .._client import DBOSClient
 from .._dbos_config import (
     _app_name_to_db_name,
     _is_valid_app_name,
+    get_system_database_url,
     is_valid_database_url,
     load_config,
 )
@@ -294,13 +295,12 @@ def migrate(
     sys_db = None
     try:
         sys_db = SystemDatabase(
-            database_url=connection_string,
+            system_database_url=get_system_database_url(config),
             engine_kwargs={
                 "pool_timeout": 30,
                 "max_overflow": 0,
                 "pool_size": 2,
             },
-            sys_db_name=sys_db_name,
         )
         app_db = ApplicationDatabase(
             database_url=connection_string,

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

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

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

@@ -12,6 +12,7 @@ from flask import Flask
 
 from dbos import DBOS, DBOSClient, DBOSConfig
 from dbos._app_db import ApplicationDatabase
+from dbos._dbos_config import get_system_database_url
 from dbos._schemas.system_database import SystemSchema
 from dbos._sys_db import SystemDatabase
 
@@ -40,7 +41,7 @@ def config() -> DBOSConfig:
 def sys_db(config: DBOSConfig) -> Generator[SystemDatabase, Any, None]:
     assert config["database_url"] is not None
     sys_db = SystemDatabase(
-        database_url=config["database_url"],
+        system_database_url=f"{config['database_url']}_dbos_sys",
         engine_kwargs={
             "pool_timeout": 30,
             "max_overflow": 0,

dbos-1.8.0a5/tests/test_async_workflow_management.py

@@ -0,0 +1,264 @@
+import asyncio
+import threading
+import time
+import uuid
+from typing import List
+
+import pytest
+
+from dbos import DBOS, Queue, SetWorkflowID
+from dbos._error import DBOSWorkflowCancelledError
+from dbos._sys_db import StepInfo, WorkflowStatus
+from tests.conftest import queue_entries_are_cleaned_up
+
+
+@pytest.mark.asyncio
+async def test_cancel_workflow_async(dbos: DBOS) -> None:
+    """Test async cancel_workflow method."""
+    steps_completed = 0
+    workflow_event = threading.Event()
+    main_thread_event = threading.Event()
+    input_val = 5
+
+    @DBOS.step()
+    def step_one() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.step()
+    def step_two() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        step_one()
+        main_thread_event.set()
+        workflow_event.wait()
+        step_two()
+        return x
+
+    # Start the workflow and cancel it async
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = DBOS.start_workflow(simple_workflow, input_val)
+    main_thread_event.wait()
+    await DBOS.cancel_workflow_async(wfid)
+    workflow_event.set()
+
+    with pytest.raises(DBOSWorkflowCancelledError):
+        handle.get_result()
+    assert steps_completed == 1
+
+
+@pytest.mark.asyncio
+async def test_resume_workflow_async(dbos: DBOS) -> None:
+    """Test async resume_workflow method."""
+    steps_completed = 0
+    workflow_event = threading.Event()
+    main_thread_event = threading.Event()
+    input_val = 5
+
+    @DBOS.step()
+    def step_one() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.step()
+    def step_two() -> None:
+        nonlocal steps_completed
+        steps_completed += 1
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        step_one()
+        main_thread_event.set()
+        workflow_event.wait()
+        step_two()
+        return x
+
+    # Start the workflow and cancel it
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = DBOS.start_workflow(simple_workflow, input_val)
+    main_thread_event.wait()
+    DBOS.cancel_workflow(wfid)
+    workflow_event.set()
+
+    with pytest.raises(DBOSWorkflowCancelledError):
+        handle.get_result()
+    assert steps_completed == 1
+
+    # Resume the workflow async
+    async_handle = await DBOS.resume_workflow_async(wfid)
+    assert (await async_handle.get_result()) == input_val
+    assert steps_completed == 2
+
+
+@pytest.mark.asyncio
+async def test_restart_workflow_async(dbos: DBOS) -> None:
+    """Test async restart_workflow method."""
+    input_val = 2
+    multiplier = 5
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        return x * multiplier
+
+    # Start the workflow, let it finish, restart it async
+    handle = DBOS.start_workflow(simple_workflow, input_val)
+    assert handle.get_result() == input_val * multiplier
+
+    forked_handle = await DBOS.restart_workflow_async(handle.workflow_id)
+    assert forked_handle.workflow_id != handle.workflow_id
+    assert (await forked_handle.get_result()) == input_val * multiplier
+
+
+@pytest.mark.asyncio
+async def test_fork_workflow_async(dbos: DBOS) -> None:
+    """Test async fork_workflow method."""
+    step_one_count = 0
+    step_two_count = 0
+    step_three_count = 0
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        return step_one(x) + step_two(x) + step_three(x)
+
+    @DBOS.step()
+    def step_one(x: int) -> int:
+        nonlocal step_one_count
+        step_one_count += 1
+        return x + 1
+
+    @DBOS.step()
+    def step_two(x: int) -> int:
+        nonlocal step_two_count
+        step_two_count += 1
+        return x + 2
+
+    @DBOS.step()
+    def step_three(x: int) -> int:
+        nonlocal step_three_count
+        step_three_count += 1
+        return x + 3
+
+    input_val = 1
+    output = 3 * input_val + 6
+
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        assert simple_workflow(input_val) == output
+
+    assert step_one_count == 1
+    assert step_two_count == 1
+    assert step_three_count == 1
+
+    # Fork from step 2 async
+    forked_handle = await DBOS.fork_workflow_async(wfid, 2)
+    assert forked_handle.workflow_id != wfid
+    assert (await forked_handle.get_result()) == output
+
+    # Verify step 1 didn't run again, but steps 2 and 3 did
+    assert step_one_count == 1
+    assert step_two_count == 2
+    assert step_three_count == 2
+
+
+@pytest.mark.asyncio
+async def test_list_workflows_async(dbos: DBOS) -> None:
+    """Test async list_workflows method."""
+    workflow_ids = []
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        return x
+
+    # Create a few workflows
+    for i in range(3):
+        wfid = str(uuid.uuid4())
+        workflow_ids.append(wfid)
+        with SetWorkflowID(wfid):
+            handle = DBOS.start_workflow(simple_workflow, i)
+            handle.get_result()
+
+    # List workflows async
+    workflows: List[WorkflowStatus] = await DBOS.list_workflows_async()
+
+    # Verify we have at least our workflows
+    assert len(workflows) >= 3
+
+    # Verify our workflow IDs are present
+    found_ids = {wf.workflow_id for wf in workflows}
+    for wfid in workflow_ids:
+        assert wfid in found_ids
+
+    # Test filtering by workflow_ids
+    filtered_workflows = await DBOS.list_workflows_async(workflow_ids=workflow_ids[:2])
+    assert len(filtered_workflows) == 2
+    assert all(wf.workflow_id in workflow_ids[:2] for wf in filtered_workflows)
+
+
+@pytest.mark.asyncio
+async def test_list_queued_workflows_async(dbos: DBOS) -> None:
+    """Test async list_queued_workflows method."""
+    queue = Queue("test_queue_async")
+    workflow_event = threading.Event()
+
+    @DBOS.workflow()
+    def blocking_workflow(x: int) -> int:
+        workflow_event.wait()
+        return x
+
+    # Enqueue a workflow but don't let it complete yet
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = queue.enqueue(blocking_workflow, 42)
+
+    # List queued workflows async while workflow is still running
+    queued_workflows = await DBOS.list_queued_workflows_async(
+        queue_name="test_queue_async"
+    )
+
+    # Verify our workflow is in the list
+    assert len(queued_workflows) >= 1
+    assert any(wf.workflow_id == wfid for wf in queued_workflows)
+
+    # Let the workflow complete
+    workflow_event.set()
+    handle.get_result()
+
+    assert queue_entries_are_cleaned_up(dbos)
+
+
+@pytest.mark.asyncio
+async def test_list_workflow_steps_async(dbos: DBOS) -> None:
+    """Test async list_workflow_steps method."""
+
+    @DBOS.workflow()
+    def simple_workflow(x: int) -> int:
+        step_one(x)
+        step_two(x)
+        return x
+
+    @DBOS.step()
+    def step_one(x: int) -> int:
+        return x + 1
+
+    @DBOS.step()
+    def step_two(x: int) -> int:
+        return x + 2
+
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        handle = DBOS.start_workflow(simple_workflow, 1)
+        handle.get_result()
+
+    # List workflow steps async
+    steps: List[StepInfo] = await DBOS.list_workflow_steps_async(wfid)
+
+    # Verify we have the expected steps (steps are returned as dictionaries)
+    assert len(steps) >= 2
+    step_names = {step["function_name"] for step in steps}
+    assert any("step_one" in name for name in step_names)
+    assert any("step_two" in name for name in step_names)

{dbos-1.8.0a3 → dbos-1.8.0a5}/tests/test_config.py

@@ -1135,6 +1135,22 @@ def test_configured_pool_default():
     dbos.destroy()
 
 
+def test_configured_pool_custom_url():
+    DBOS.destroy()
+    config: DBOSConfig = {
+        "name": "test-app",
+        "database_url": f"postgres://postgres:{quote(os.environ.get('PGPASSWORD', 'dbos'))}@localhost:5432/postgres",
+        "system_database_url": f"postgres://postgres:{quote(os.environ.get('PGPASSWORD', 'dbos'))}@localhost:5432/dbostesturl",
+    }
+
+    dbos = DBOS(config=config)
+    dbos.launch()
+    assert "postgres" in dbos._app_db.engine.url
+    assert "dbostesturl" in dbos._sys_db.engine.url
+
+    dbos.destroy()
+
+
 def test_configured_pool_user_provided():
     DBOS.destroy()
     config: DBOSConfig = {