dbos 1.8.0a1__tar.gz → 1.8.0a5__tar.gz

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

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

{dbos-1.8.0a1 → 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.0a1 → dbos-1.8.0a5}/dbos/_context.py

@@ -10,7 +10,7 @@ from enum import Enum
 from types import TracebackType
 from typing import List, Literal, Optional, Type, TypedDict
 
-from opentelemetry.trace import Span, Status, StatusCode
+from opentelemetry.trace import Span, Status, StatusCode, use_span
 from sqlalchemy.orm import Session
 
 from dbos._utils import GlobalParams
@@ -68,6 +68,20 @@ class StepStatus:
     max_attempts: Optional[int]
 
 
+@dataclass
+class ContextSpan:
+    """
+    A span that is used to track the context of a workflow or step execution.
+
+    Attributes:
+        span: The OpenTelemetry span object.
+        context_manager: The context manager that is used to manage the span's lifecycle.
+    """
+
+    span: Span
+    context_manager: AbstractContextManager[Span]
+
+
 class DBOSContext:
     def __init__(self) -> None:
         self.executor_id = GlobalParams.executor_id
@@ -86,7 +100,7 @@ class DBOSContext:
         self.curr_step_function_id: int = -1
         self.curr_tx_function_id: int = -1
         self.sql_session: Optional[Session] = None
-        self.spans: list[Span] = []
+        self.context_spans: list[ContextSpan] = []
 
         self.authenticated_user: Optional[str] = None
         self.authenticated_roles: Optional[List[str]] = None
@@ -202,8 +216,8 @@ class DBOSContext:
         self._end_span(exc_value)
 
     def get_current_span(self) -> Optional[Span]:
-        if len(self.spans):
-            return self.spans[-1]
+        if len(self.context_spans) > 0:
+            return self.context_spans[-1].span
         return None
 
     def _start_span(self, attributes: TracedAttributes) -> None:
@@ -218,27 +232,38 @@ class DBOSContext:
         )
         attributes["authenticatedUserAssumedRole"] = self.assumed_role
         span = dbos_tracer.start_span(
-            attributes, parent=self.spans[-1] if len(self.spans) > 0 else None
+            attributes,
+            parent=self.context_spans[-1].span if len(self.context_spans) > 0 else None,
+        )
+        # Activate the current span
+        cm = use_span(
+            span,
+            end_on_exit=False,
+            record_exception=False,
+            set_status_on_exception=False,
         )
-        self.spans.append(span)
+        self.context_spans.append(ContextSpan(span, cm))
+        cm.__enter__()
 
     def _end_span(self, exc_value: Optional[BaseException]) -> None:
+        context_span = self.context_spans.pop()
         if exc_value is None:
-            self.spans[-1].set_status(Status(StatusCode.OK))
+            context_span.span.set_status(Status(StatusCode.OK))
         else:
-            self.spans[-1].set_status(
+            context_span.span.set_status(
                 Status(StatusCode.ERROR, description=str(exc_value))
             )
-        dbos_tracer.end_span(self.spans.pop())
+        dbos_tracer.end_span(context_span.span)
+        context_span.context_manager.__exit__(None, None, None)
 
     def set_authentication(
         self, user: Optional[str], roles: Optional[List[str]]
     ) -> None:
         self.authenticated_user = user
         self.authenticated_roles = roles
-        if user is not None and len(self.spans) > 0:
-            self.spans[-1].set_attribute("authenticatedUser", user)
-            self.spans[-1].set_attribute(
+        if user is not None and len(self.context_spans) > 0:
+            self.context_spans[-1].span.set_attribute("authenticatedUser", user)
+            self.context_spans[-1].span.set_attribute(
                 "authenticatedUserRoles", json.dumps(roles) if roles is not None else ""
             )
 

{dbos-1.8.0a1 → 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.0a1 → 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.0a1 → dbos-1.8.0a5}/dbos/_error.py

@@ -126,7 +126,7 @@ class DBOSDeadLetterQueueError(DBOSException):
 
     def __init__(self, wf_id: str, max_retries: int):
         super().__init__(
-            f"Workflow {wf_id} has been moved to the dead-letter queue after exceeding the maximum of ${max_retries} retries",
+            f"Workflow {wf_id} has been moved to the dead-letter queue after exceeding the maximum of {max_retries} retries",
             dbos_error_code=DBOSErrorCode.DeadLetterQueueError.value,
         )
 

{dbos-1.8.0a1 → 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
@@ -437,7 +430,14 @@ class SystemDatabase:
 
         # Values to update when a row already exists for this workflow
         update_values: dict[str, Any] = {
-            "recovery_attempts": SystemSchema.workflow_status.c.recovery_attempts + 1,
+            "recovery_attempts": sa.case(
+                (
+                    SystemSchema.workflow_status.c.status
+                    != WorkflowStatusString.ENQUEUED.value,
+                    SystemSchema.workflow_status.c.recovery_attempts + 1,
+                ),
+                else_=SystemSchema.workflow_status.c.recovery_attempts,
+            ),
             "updated_at": func.extract("epoch", func.now()) * 1000,
         }
         # Don't update an existing executor ID when enqueueing a workflow.

{dbos-1.8.0a1 → 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.0a1 → dbos-1.8.0a5}/pyproject.toml

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

{dbos-1.8.0a1 → 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,