dbos 0.24.1__py3-none-any.whl → 0.25.0__py3-none-any.whl

dbos/_dbos.py

@@ -11,7 +11,6 @@ import threading
 import traceback
 import uuid
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
 from logging import Logger
 from typing import (
     TYPE_CHECKING,
@@ -28,17 +27,23 @@ from typing import (
     TypeVar,
     Union,
     cast,
-    overload,
 )
 
 from opentelemetry.trace import Span
 
+from dbos import _serialization
 from dbos._conductor.conductor import ConductorWebsocket
 from dbos._utils import GlobalParams
+from dbos._workflow_commands import (
+    WorkflowStatus,
+    list_queued_workflows,
+    list_workflows,
+)
 
 from ._classproperty import classproperty
 from ._core import (
     TEMP_SEND_WF_NAME,
+    WorkflowHandleAsyncPolling,
     WorkflowHandlePolling,
     decorate_step,
     decorate_transaction,
@@ -85,6 +90,7 @@ from ._admin_server import AdminServer
 from ._app_db import ApplicationDatabase
 from ._context import (
     EnterDBOSStep,
+    StepStatus,
     TracedAttributes,
     assert_current_dbos_context,
     get_local_dbos_context,
@@ -107,6 +113,7 @@ from ._error import (
 )
 from ._logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
 from ._sys_db import SystemDatabase
+from ._workflow_commands import WorkflowStatus, get_workflow
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F
 # There are cases where the parameters P and return value R should be separate
@@ -119,17 +126,6 @@ R = TypeVar("R", covariant=True)  # A generic type for workflow return values
 
 T = TypeVar("T")
 
-
-class DBOSCallProtocol(Protocol[P, R]):
-    __name__: str
-    __qualname__: str
-
-    def __call__(*args: P.args, **kwargs: P.kwargs) -> R: ...
-
-
-Workflow: TypeAlias = DBOSCallProtocol[P, R]
-
-
 IsolationLevel = Literal[
     "SERIALIZABLE",
     "REPEATABLE READ",
@@ -162,7 +158,7 @@ RegisteredJob = Tuple[
 
 class DBOSRegistry:
     def __init__(self) -> None:
-        self.workflow_info_map: dict[str, Workflow[..., Any]] = {}
+        self.workflow_info_map: dict[str, Callable[..., Any]] = {}
         self.function_type_map: dict[str, str] = {}
         self.class_info_map: dict[str, type] = {}
         self.instance_info_map: dict[str, object] = {}
@@ -449,8 +445,12 @@ class DBOS:
             dbos_logger.info(f"Executor ID: {GlobalParams.executor_id}")
             dbos_logger.info(f"Application version: {GlobalParams.app_version}")
             self._executor_field = ThreadPoolExecutor(max_workers=64)
-            self._sys_db_field = SystemDatabase(self.config, debug_mode=debug_mode)
-            self._app_db_field = ApplicationDatabase(self.config, debug_mode=debug_mode)
+            self._sys_db_field = SystemDatabase(
+                self.config["database"], debug_mode=debug_mode
+            )
+            self._app_db_field = ApplicationDatabase(
+                self.config["database"], debug_mode=debug_mode
+            )
 
             if debug_mode:
                 return
@@ -489,14 +489,6 @@ class DBOS:
             notification_listener_thread.start()
             self._background_threads.append(notification_listener_thread)
 
-            # Start flush workflow buffers thread
-            flush_workflow_buffers_thread = threading.Thread(
-                target=self._sys_db.flush_workflow_buffers,
-                daemon=True,
-            )
-            flush_workflow_buffers_thread.start()
-            self._background_threads.append(flush_workflow_buffers_thread)
-
             # Start the queue thread
             evt = threading.Event()
             self.stop_events.append(evt)
@@ -552,6 +544,10 @@ class DBOS:
         """
         if _dbos_global_instance is not None:
             _dbos_global_instance._reset_system_database()
+        else:
+            dbos_logger.warning(
+                "reset_system_database has no effect because global DBOS object does not exist"
+            )
 
     def _reset_system_database(self) -> None:
         assert (
@@ -706,7 +702,7 @@ class DBOS:
     @classmethod
     def start_workflow(
         cls,
-        func: Workflow[P, R],
+        func: Callable[P, R],
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> WorkflowHandle[R]:
@@ -716,7 +712,7 @@ class DBOS:
     @classmethod
     async def start_workflow_async(
         cls,
-        func: Workflow[P, Coroutine[Any, Any, R]],
+        func: Callable[P, Coroutine[Any, Any, R]],
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> WorkflowHandleAsync[R]:
@@ -728,34 +724,39 @@ class DBOS:
     @classmethod
     def get_workflow_status(cls, workflow_id: str) -> Optional[WorkflowStatus]:
         """Return the status of a workflow execution."""
+        sys_db = _get_dbos_instance()._sys_db
         ctx = get_local_dbos_context()
         if ctx and ctx.is_within_workflow():
             ctx.function_id += 1
-            stat = _get_dbos_instance()._sys_db.get_workflow_status_within_wf(
-                workflow_id, ctx.workflow_id, ctx.function_id
+            res = sys_db.check_operation_execution(ctx.workflow_id, ctx.function_id)
+            if res is not None:
+                if res["output"]:
+                    resstat: WorkflowStatus = _serialization.deserialize(res["output"])
+                    return resstat
+                else:
+                    raise DBOSException(
+                        "Workflow status record not found. This should not happen! \033[1m Hint: Check if your workflow is deterministic.\033[0m"
+                    )
+        stat = get_workflow(_get_dbos_instance()._sys_db, workflow_id, True)
+
+        if ctx and ctx.is_within_workflow():
+            sys_db.record_operation_result(
+                {
+                    "workflow_uuid": ctx.workflow_id,
+                    "function_id": ctx.function_id,
+                    "function_name": "DBOS.getStatus",
+                    "output": _serialization.serialize(stat),
+                    "error": None,
+                }
             )
-        else:
-            stat = _get_dbos_instance()._sys_db.get_workflow_status(workflow_id)
-        if stat is None:
-            return None
-
-        return WorkflowStatus(
-            workflow_id=workflow_id,
-            status=stat["status"],
-            name=stat["name"],
-            executor_id=stat["executor_id"],
-            recovery_attempts=stat["recovery_attempts"],
-            class_name=stat["class_name"],
-            config_name=stat["config_name"],
-            queue_name=stat["queue_name"],
-            authenticated_user=stat["authenticated_user"],
-            assumed_role=stat["assumed_role"],
-            authenticated_roles=(
-                json.loads(stat["authenticated_roles"])
-                if stat["authenticated_roles"] is not None
-                else None
-            ),
-        )
+        return stat
+
+    @classmethod
+    async def get_workflow_status_async(
+        cls, workflow_id: str
+    ) -> Optional[WorkflowStatus]:
+        """Return the status of a workflow execution."""
+        return await asyncio.to_thread(cls.get_workflow_status, workflow_id)
 
     @classmethod
     def retrieve_workflow(
@@ -769,6 +770,18 @@ class DBOS:
                 raise DBOSNonExistentWorkflowError(workflow_id)
         return WorkflowHandlePolling(workflow_id, dbos)
 
+    @classmethod
+    async def retrieve_workflow_async(
+        cls, workflow_id: str, existing_workflow: bool = True
+    ) -> WorkflowHandleAsync[R]:
+        """Return a `WorkflowHandle` for a workflow execution."""
+        dbos = _get_dbos_instance()
+        if existing_workflow:
+            stat = await dbos.get_workflow_status_async(workflow_id)
+            if stat is None:
+                raise DBOSNonExistentWorkflowError(workflow_id)
+        return WorkflowHandleAsyncPolling(workflow_id, dbos)
+
     @classmethod
     def send(
         cls, destination_id: str, message: Any, topic: Optional[str] = None
@@ -943,6 +956,60 @@ class DBOS:
         _get_or_create_dbos_registry().clear_workflow_cancelled(workflow_id)
         return execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
 
+    @classmethod
+    def list_workflows(
+        cls,
+        *,
+        workflow_ids: Optional[List[str]] = None,
+        status: Optional[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,
+    ) -> List[WorkflowStatus]:
+        return list_workflows(
+            _get_dbos_instance()._sys_db,
+            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,
+        )
+
+    @classmethod
+    def list_queued_workflows(
+        cls,
+        *,
+        queue_name: Optional[str] = None,
+        status: Optional[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,
+    ) -> List[WorkflowStatus]:
+        return list_queued_workflows(
+            _get_dbos_instance()._sys_db,
+            queue_name=queue_name,
+            status=status,
+            start_time=start_time,
+            end_time=end_time,
+            name=name,
+            limit=limit,
+            offset=offset,
+            sort_desc=sort_desc,
+        )
+
     @classproperty
     def logger(cls) -> Logger:
         """Return the DBOS `Logger` for the current context."""
@@ -981,6 +1048,23 @@ class DBOS:
         ), "workflow_id is only available within a DBOS operation."
         return ctx.workflow_id
 
+    @classproperty
+    def step_id(cls) -> int:
+        """Return the step ID for the current context. This is a unique identifier of the current step within the workflow."""
+        ctx = assert_current_dbos_context()
+        assert (
+            ctx.is_within_workflow()
+        ), "step_id is only available within a DBOS workflow."
+        return ctx.function_id
+
+    @classproperty
+    def step_status(cls) -> StepStatus:
+        """Return the status of the currently executing step."""
+        ctx = assert_current_dbos_context()
+        assert ctx.is_step(), "step_status is only available within a DBOS step."
+        assert ctx.step_status is not None
+        return ctx.step_status
+
     @classproperty
     def parent_workflow_id(cls) -> str:
         """
@@ -1035,41 +1119,6 @@ class DBOS:
         ctx.authenticated_roles = authenticated_roles
 
 
-@dataclass
-class WorkflowStatus:
-    """
-    Status of workflow execution.
-
-    This captures the state of a workflow execution at a point in time.
-
-    Attributes:
-        workflow_id(str):  The ID of the workflow execution
-        status(str):  The status of the execution, from `WorkflowStatusString`
-        name(str): The workflow function name
-        executor_id(str): The ID of the executor running the workflow
-        class_name(str): For member functions, the name of the class containing the workflow function
-        config_name(str): For instance member functions, the name of the class instance for the execution
-        queue_name(str): For workflows that are or were queued, the queue name
-        authenticated_user(str): The user who invoked the workflow
-        assumed_role(str): The access role used by the user to allow access to the workflow function
-        authenticated_roles(List[str]): List of all access roles available to the authenticated user
-        recovery_attempts(int): Number of times the workflow has been restarted (usually by recovery)
-
-    """
-
-    workflow_id: str
-    status: str
-    name: str
-    executor_id: Optional[str]
-    class_name: Optional[str]
-    config_name: Optional[str]
-    queue_name: Optional[str]
-    authenticated_user: Optional[str]
-    assumed_role: Optional[str]
-    authenticated_roles: Optional[List[str]]
-    recovery_attempts: Optional[int]
-
-
 class WorkflowHandle(Generic[R], Protocol):
     """
     Handle to a workflow function.

dbos/_fastapi.py

@@ -63,7 +63,10 @@ class LifespanMiddleware:
         if scope["type"] == "lifespan":
 
             async def wrapped_send(message: MutableMapping[str, Any]) -> None:
-                if message["type"] == "lifespan.startup.complete":
+                if (
+                    message["type"] == "lifespan.startup.complete"
+                    and not self.dbos._launched
+                ):
                     self.dbos._launch()
                 elif message["type"] == "lifespan.shutdown.complete":
                     self.dbos._destroy()

dbos/_logger.py

@@ -7,6 +7,7 @@ from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
 from opentelemetry.sdk.resources import Resource
+from opentelemetry.trace.span import format_trace_id
 
 from dbos._utils import GlobalParams
 
@@ -26,6 +27,19 @@ class DBOSLogTransformer(logging.Filter):
         record.applicationID = self.app_id
         record.applicationVersion = GlobalParams.app_version
         record.executorID = GlobalParams.executor_id
+
+        # If available, decorate the log entry with Workflow ID and Trace ID
+        from dbos._context import get_local_dbos_context
+
+        ctx = get_local_dbos_context()
+        if ctx:
+            if ctx.is_within_workflow():
+                record.operationUUID = ctx.workflow_id
+            span = ctx.get_current_span()
+            if span:
+                trace_id = format_trace_id(span.get_span_context().trace_id)
+                record.traceId = trace_id
+
         return True
 
 

dbos/_migrations/versions/f4b9b32ba814_functionname_childid_op_outputs.py

@@ -0,0 +1,46 @@
+"""functionname_childid_op_outputs
+
+Revision ID: f4b9b32ba814
+Revises: 04ca4f231047
+Create Date: 2025-03-21 14:32:43.091074
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "f4b9b32ba814"
+down_revision: Union[str, None] = "04ca4f231047"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "operation_outputs",
+        sa.Column(
+            "function_name",
+            sa.Text(),
+            nullable=False,
+            server_default="",
+        ),
+        schema="dbos",
+    )
+
+    op.add_column(
+        "operation_outputs",
+        sa.Column(
+            "child_workflow_id",
+            sa.Text(),
+            nullable=True,
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("operation_outputs", "function_name", schema="dbos")
+    op.drop_column("operation_outputs", "child_workflow_id", schema="dbos")

dbos/_outcome.py

@@ -4,6 +4,8 @@ import inspect
 import time
 from typing import Any, Callable, Coroutine, Optional, Protocol, TypeVar, Union, cast
 
+from dbos._context import EnterDBOSStepRetry
+
 T = TypeVar("T")
 R = TypeVar("R")
 
@@ -98,7 +100,8 @@ class Immediate(Outcome[T]):
     ) -> T:
         for i in range(attempts):
             try:
-                return func()
+                with EnterDBOSStepRetry(i, attempts):
+                    return func()
             except Exception as exp:
                 wait_time = on_exception(i, exp)
                 time.sleep(wait_time)
@@ -184,7 +187,8 @@ class Pending(Outcome[T]):
     ) -> T:
         for i in range(attempts):
             try:
-                return await func()
+                with EnterDBOSStepRetry(i, attempts):
+                    return await func()
             except Exception as exp:
                 wait_time = on_exception(i, exp)
                 await asyncio.sleep(wait_time)

dbos/_queue.py

@@ -1,6 +1,6 @@
 import threading
 import traceback
-from typing import TYPE_CHECKING, Any, Coroutine, Optional, TypedDict
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypedDict
 
 from psycopg import errors
 from sqlalchemy.exc import OperationalError
@@ -10,7 +10,7 @@ from dbos._utils import GlobalParams
 from ._core import P, R, execute_workflow_by_id, start_workflow, start_workflow_async
 
 if TYPE_CHECKING:
-    from ._dbos import DBOS, Workflow, WorkflowHandle, WorkflowHandleAsync
+    from ._dbos import DBOS, WorkflowHandle, WorkflowHandleAsync
 
 
 class QueueRateLimit(TypedDict):
@@ -59,7 +59,7 @@ class Queue:
         registry.queue_info_map[self.name] = self
 
     def enqueue(
-        self, func: "Workflow[P, R]", *args: P.args, **kwargs: P.kwargs
+        self, func: "Callable[P, R]", *args: P.args, **kwargs: P.kwargs
     ) -> "WorkflowHandle[R]":
         from ._dbos import _get_dbos_instance
 
@@ -68,7 +68,7 @@ class Queue:
 
     async def enqueue_async(
         self,
-        func: "Workflow[P, Coroutine[Any, Any, R]]",
+        func: "Callable[P, Coroutine[Any, Any, R]]",
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> "WorkflowHandleAsync[R]":
@@ -85,7 +85,7 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
         for _, queue in dbos._registry.queue_info_map.items():
             try:
                 wf_ids = dbos._sys_db.start_queued_workflows(
-                    queue, GlobalParams.executor_id
+                    queue, GlobalParams.executor_id, GlobalParams.app_version
                 )
                 for id in wf_ids:
                     execute_workflow_by_id(dbos, id)

dbos/_schemas/system_database.py

@@ -71,8 +71,10 @@ class SystemSchema:
             nullable=False,
         ),
         Column("function_id", Integer, nullable=False),
+        Column("function_name", Text, nullable=False, default=""),
         Column("output", Text, nullable=True),
         Column("error", Text, nullable=True),
+        Column("child_workflow_id", Text, nullable=True),
         PrimaryKeyConstraint("workflow_uuid", "function_id"),
     )