dbos 0.24.0a15__py3-none-any.whl → 0.25.0__py3-none-any.whl

dbos/__init__.py

@@ -1,18 +1,22 @@
 from . import _error as error
+from ._client import DBOSClient, EnqueueOptions
 from ._context import DBOSContextEnsure, DBOSContextSetAuth, SetWorkflowID
-from ._dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
+from ._dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle
 from ._dbos_config import ConfigFile, DBOSConfig, get_dbos_database_url, load_config
 from ._kafka_message import KafkaMessage
 from ._queue import Queue
 from ._sys_db import GetWorkflowsInput, WorkflowStatusString
+from ._workflow_commands import WorkflowStatus
 
 __all__ = [
     "ConfigFile",
     "DBOSConfig",
     "DBOS",
+    "DBOSClient",
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",
     "DBOSContextSetAuth",
+    "EnqueueOptions",
     "GetWorkflowsInput",
     "KafkaMessage",
     "SetWorkflowID",

dbos/__main__.py

@@ -4,6 +4,9 @@ from typing import NoReturn, Optional, Union
 
 from dbos.cli.cli import app
 
+# This is used by the debugger to execute DBOS as a module.
+# Never used otherwise.
+
 
 def main() -> NoReturn:
     # Modify sys.argv[0] to remove script or executable extensions

dbos/_admin_server.py

@@ -20,6 +20,7 @@ _workflow_queues_metadata_path = "/dbos-workflow-queues-metadata"
 # /workflows/:workflow_id/cancel
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
+# /workflows/:workflow_id/steps
 
 
 class AdminServer:
@@ -86,8 +87,16 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self._end_headers()
             self.wfile.write(json.dumps(queue_metadata_array).encode("utf-8"))
         else:
-            self.send_response(404)
-            self._end_headers()
+            steps_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/steps$", self.path
+            )
+
+            if steps_match:
+                workflow_id = steps_match.group("workflow_id")
+                self._handle_steps(workflow_id)
+            else:
+                self.send_response(404)
+                self._end_headers()
 
     def do_POST(self) -> None:
         content_length = int(
@@ -149,6 +158,23 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         self.send_response(204)
         self._end_headers()
 
+    def _handle_steps(self, workflow_id: str) -> None:
+        steps = self.dbos._sys_db.get_workflow_steps(workflow_id)
+
+        updated_steps = [
+            {
+                **step,
+                "output": str(step["output"]) if step["output"] is not None else None,
+                "error": str(step["error"]) if step["error"] is not None else None,
+            }
+            for step in steps
+        ]
+
+        json_steps = json.dumps(updated_steps).encode("utf-8")
+        self.send_response(200)
+        self._end_headers()
+        self.wfile.write(json_steps)
+
 
 # Be consistent with DBOS-TS response.
 class PerfUtilization(TypedDict):

dbos/_app_db.py

@@ -27,19 +27,18 @@ class RecordedResult(TypedDict):
 
 class ApplicationDatabase:
 
-    def __init__(self, config: ConfigFile, *, debug_mode: bool = False):
-        self.config = config
+    def __init__(self, database: DatabaseConfig, *, debug_mode: bool = False):
 
-        app_db_name = config["database"]["app_db_name"]
+        app_db_name = database["app_db_name"]
 
         # If the application database does not already exist, create it
         if not debug_mode:
             postgres_db_url = sa.URL.create(
                 "postgresql+psycopg",
-                username=config["database"]["username"],
-                password=config["database"]["password"],
-                host=config["database"]["hostname"],
-                port=config["database"]["port"],
+                username=database["username"],
+                password=database["password"],
+                host=database["hostname"],
+                port=database["port"],
                 database="postgres",
             )
             postgres_db_engine = sa.create_engine(postgres_db_url)
@@ -55,25 +54,25 @@ class ApplicationDatabase:
         # Create a connection pool for the application database
         app_db_url = sa.URL.create(
             "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
+            username=database["username"],
+            password=database["password"],
+            host=database["hostname"],
+            port=database["port"],
             database=app_db_name,
         )
 
         connect_args = {}
         if (
-            "connectionTimeoutMillis" in config["database"]
-            and config["database"]["connectionTimeoutMillis"]
+            "connectionTimeoutMillis" in database
+            and database["connectionTimeoutMillis"]
         ):
             connect_args["connect_timeout"] = int(
-                config["database"]["connectionTimeoutMillis"] / 1000
+                database["connectionTimeoutMillis"] / 1000
             )
 
         self.engine = sa.create_engine(
             app_db_url,
-            pool_size=config["database"]["app_db_pool_size"],
+            pool_size=database["app_db_pool_size"],
             max_overflow=0,
             pool_timeout=30,
             connect_args=connect_args,

dbos/_client.py

@@ -0,0 +1,206 @@
+import asyncio
+import sys
+import uuid
+from typing import Any, Generic, Optional, TypedDict, TypeVar
+
+if sys.version_info < (3, 11):
+    from typing_extensions import NotRequired
+else:
+    from typing import NotRequired
+
+from dbos import _serialization
+from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
+from dbos._dbos_config import parse_database_url_to_dbconfig
+from dbos._error import DBOSNonExistentWorkflowError
+from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
+from dbos._serialization import WorkflowInputs
+from dbos._sys_db import SystemDatabase, WorkflowStatusInternal, WorkflowStatusString
+from dbos._workflow_commands import WorkflowStatus, get_workflow
+
+R = TypeVar("R", covariant=True)  # A generic type for workflow return values
+
+
+class EnqueueOptions(TypedDict):
+    workflow_name: str
+    workflow_class_name: NotRequired[str]
+    queue_name: str
+    app_version: NotRequired[str]
+    workflow_id: NotRequired[str]
+
+
+class WorkflowHandleClientPolling(Generic[R]):
+
+    def __init__(self, workflow_id: str, sys_db: SystemDatabase):
+        self.workflow_id = workflow_id
+        self._sys_db = sys_db
+
+    def get_workflow_id(self) -> str:
+        return self.workflow_id
+
+    def get_result(self) -> R:
+        res: R = self._sys_db.await_workflow_result(self.workflow_id)
+        return res
+
+    def get_status(self) -> "WorkflowStatus":
+        status = get_workflow(self._sys_db, self.workflow_id, True)
+        if status is None:
+            raise DBOSNonExistentWorkflowError(self.workflow_id)
+        return status
+
+
+class WorkflowHandleClientAsyncPolling(Generic[R]):
+
+    def __init__(self, workflow_id: str, sys_db: SystemDatabase):
+        self.workflow_id = workflow_id
+        self._sys_db = sys_db
+
+    def get_workflow_id(self) -> str:
+        return self.workflow_id
+
+    async def get_result(self) -> R:
+        res: R = await asyncio.to_thread(
+            self._sys_db.await_workflow_result, self.workflow_id
+        )
+        return res
+
+    async def get_status(self) -> "WorkflowStatus":
+        status = await asyncio.to_thread(
+            get_workflow, self._sys_db, self.workflow_id, True
+        )
+        if status is None:
+            raise DBOSNonExistentWorkflowError(self.workflow_id)
+        return status
+
+
+class DBOSClient:
+    def __init__(self, database_url: str, *, system_database: Optional[str] = None):
+        db_config = parse_database_url_to_dbconfig(database_url)
+        if system_database is not None:
+            db_config["sys_db_name"] = system_database
+        self._sys_db = SystemDatabase(db_config)
+
+    def destroy(self) -> None:
+        self._sys_db.destroy()
+
+    def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
+        workflow_name = options["workflow_name"]
+        queue_name = options["queue_name"]
+
+        workflow_class_name = options.get("workflow_class_name")
+        app_version = options.get("app_version")
+        max_recovery_attempts = options.get("max_recovery_attempts")
+        if max_recovery_attempts is None:
+            max_recovery_attempts = DEFAULT_MAX_RECOVERY_ATTEMPTS
+        workflow_id = options.get("workflow_id")
+        if workflow_id is None:
+            workflow_id = str(uuid.uuid4())
+
+        status: WorkflowStatusInternal = {
+            "workflow_uuid": workflow_id,
+            "status": WorkflowStatusString.ENQUEUED.value,
+            "name": workflow_name,
+            "class_name": workflow_class_name,
+            "queue_name": queue_name,
+            "app_version": app_version,
+            "config_name": None,
+            "authenticated_user": None,
+            "assumed_role": None,
+            "authenticated_roles": None,
+            "request": None,
+            "output": None,
+            "error": None,
+            "created_at": None,
+            "updated_at": None,
+            "executor_id": None,
+            "recovery_attempts": None,
+            "app_id": None,
+        }
+
+        inputs: WorkflowInputs = {
+            "args": args,
+            "kwargs": kwargs,
+        }
+
+        wf_status = self._sys_db.insert_workflow_status(status)
+        self._sys_db.update_workflow_inputs(
+            workflow_id, _serialization.serialize_args(inputs)
+        )
+        if wf_status == WorkflowStatusString.ENQUEUED.value:
+            self._sys_db.enqueue(workflow_id, queue_name)
+        return workflow_id
+
+    def enqueue(
+        self, options: EnqueueOptions, *args: Any, **kwargs: Any
+    ) -> WorkflowHandle[R]:
+        workflow_id = self._enqueue(options, *args, **kwargs)
+        return WorkflowHandleClientPolling[R](workflow_id, self._sys_db)
+
+    async def enqueue_async(
+        self, options: EnqueueOptions, *args: Any, **kwargs: Any
+    ) -> WorkflowHandleAsync[R]:
+        workflow_id = await asyncio.to_thread(self._enqueue, options, *args, **kwargs)
+        return WorkflowHandleClientAsyncPolling[R](workflow_id, self._sys_db)
+
+    def retrieve_workflow(self, workflow_id: str) -> WorkflowHandle[R]:
+        status = get_workflow(self._sys_db, workflow_id, True)
+        if status is None:
+            raise DBOSNonExistentWorkflowError(workflow_id)
+        return WorkflowHandleClientPolling[R](workflow_id, self._sys_db)
+
+    async def retrieve_workflow_async(self, workflow_id: str) -> WorkflowHandleAsync[R]:
+        status = asyncio.to_thread(get_workflow, self._sys_db, workflow_id, True)
+        if status is None:
+            raise DBOSNonExistentWorkflowError(workflow_id)
+        return WorkflowHandleClientAsyncPolling[R](workflow_id, self._sys_db)
+
+    def send(
+        self,
+        destination_id: str,
+        message: Any,
+        topic: Optional[str] = None,
+        idempotency_key: Optional[str] = None,
+    ) -> None:
+        idempotency_key = idempotency_key if idempotency_key else str(uuid.uuid4())
+        status: WorkflowStatusInternal = {
+            "workflow_uuid": f"{destination_id}-{idempotency_key}",
+            "status": WorkflowStatusString.SUCCESS.value,
+            "name": "temp_workflow-send-client",
+            "class_name": None,
+            "queue_name": None,
+            "config_name": None,
+            "authenticated_user": None,
+            "assumed_role": None,
+            "authenticated_roles": None,
+            "request": None,
+            "output": None,
+            "error": None,
+            "created_at": None,
+            "updated_at": None,
+            "executor_id": None,
+            "recovery_attempts": None,
+            "app_id": None,
+            "app_version": None,
+        }
+        self._sys_db.insert_workflow_status(status)
+        self._sys_db.send(status["workflow_uuid"], 0, destination_id, message, topic)
+
+    async def send_async(
+        self,
+        destination_id: str,
+        message: Any,
+        topic: Optional[str] = None,
+        idempotency_key: Optional[str] = None,
+    ) -> None:
+        return await asyncio.to_thread(
+            self.send, destination_id, message, topic, idempotency_key
+        )
+
+    def get_event(self, workflow_id: str, key: str, timeout_seconds: float = 60) -> Any:
+        return self._sys_db.get_event(workflow_id, key, timeout_seconds)
+
+    async def get_event_async(
+        self, workflow_id: str, key: str, timeout_seconds: float = 60
+    ) -> Any:
+        return await asyncio.to_thread(
+            self.get_event, workflow_id, key, timeout_seconds
+        )

dbos/_conductor/conductor.py

@@ -9,7 +9,12 @@ from websockets.sync.client import connect
 from websockets.sync.connection import Connection
 
 from dbos._utils import GlobalParams
-from dbos._workflow_commands import get_workflow, list_queued_workflows, list_workflows
+from dbos._workflow_commands import (
+    get_workflow,
+    list_queued_workflows,
+    list_workflow_steps,
+    list_workflows,
+)
 
 from . import protocol as p
 
@@ -203,7 +208,7 @@ class ConductorWebsocket(threading.Thread):
                                 info = get_workflow(
                                     self.dbos._sys_db,
                                     get_workflow_message.workflow_id,
-                                    getRequest=False,
+                                    get_request=False,
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when getting workflow {get_workflow_message.workflow_id}: {traceback.format_exc()}"
@@ -243,6 +248,32 @@ class ConductorWebsocket(threading.Thread):
                                 )
                             )
                             websocket.send(exist_pending_workflows_response.to_json())
+                        elif msg_type == p.MessageType.LIST_STEPS:
+                            list_steps_message = p.ListStepsRequest.from_json(message)
+                            step_info = None
+                            try:
+                                step_info = list_workflow_steps(
+                                    self.dbos._sys_db,
+                                    list_steps_message.workflow_id,
+                                )
+                            except Exception as e:
+                                error_message = f"Exception encountered when getting workflow {list_steps_message.workflow_id}: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+
+                            list_steps_response = p.ListStepsResponse(
+                                type=p.MessageType.LIST_STEPS,
+                                request_id=base_message.request_id,
+                                output=(
+                                    [
+                                        p.WorkflowSteps.from_step_info(i)
+                                        for i in step_info
+                                    ]
+                                    if step_info is not None
+                                    else None
+                                ),
+                                error_message=error_message,
+                            )
+                            websocket.send(list_steps_response.to_json())
                         else:
                             self.dbos.logger.warning(
                                 f"Unexpected message type: {msg_type}"

dbos/_conductor/protocol.py

@@ -3,7 +3,8 @@ from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import List, Optional, Type, TypedDict, TypeVar
 
-from dbos._workflow_commands import WorkflowInformation
+from dbos._sys_db import StepInfo
+from dbos._workflow_commands import WorkflowStatus
 
 
 class MessageType(str, Enum):
@@ -16,6 +17,7 @@ class MessageType(str, Enum):
     RESTART = "restart"
     GET_WORKFLOW = "get_workflow"
     EXIST_PENDING_WORKFLOWS = "exist_pending_workflows"
+    LIST_STEPS = "list_steps"
 
 
 T = TypeVar("T", bound="BaseMessage")
@@ -141,27 +143,33 @@ class WorkflowsOutput:
     ExecutorID: Optional[str]
 
     @classmethod
-    def from_workflow_information(cls, info: WorkflowInformation) -> "WorkflowsOutput":
+    def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
         # Convert fields to strings as needed
         created_at_str = str(info.created_at) if info.created_at is not None else None
         updated_at_str = str(info.updated_at) if info.updated_at is not None else None
         inputs_str = str(info.input) if info.input is not None else None
         outputs_str = str(info.output) if info.output is not None else None
+        error_str = str(info.error) if info.error is not None else None
         request_str = str(info.request) if info.request is not None else None
+        roles_str = (
+            str(info.authenticated_roles)
+            if info.authenticated_roles is not None
+            else None
+        )
 
         return cls(
             WorkflowUUID=info.workflow_id,
             Status=info.status,
-            WorkflowName=info.workflow_name,
-            WorkflowClassName=info.workflow_class_name,
-            WorkflowConfigName=info.workflow_config_name,
+            WorkflowName=info.name,
+            WorkflowClassName=info.class_name,
+            WorkflowConfigName=info.config_name,
             AuthenticatedUser=info.authenticated_user,
             AssumedRole=info.assumed_role,
-            AuthenticatedRoles=info.authenticated_roles,
+            AuthenticatedRoles=roles_str,
             Input=inputs_str,
             Output=outputs_str,
             Request=request_str,
-            Error=info.error,
+            Error=error_str,
             CreatedAt=created_at_str,
             UpdatedAt=updated_at_str,
             QueueName=info.queue_name,
@@ -170,6 +178,27 @@ class WorkflowsOutput:
         )
 
 
+@dataclass
+class WorkflowSteps:
+    function_id: int
+    function_name: str
+    output: Optional[str]
+    error: Optional[str]
+    child_workflow_id: Optional[str]
+
+    @classmethod
+    def from_step_info(cls, info: StepInfo) -> "WorkflowSteps":
+        output_str = str(info["output"]) if info["output"] is not None else None
+        error_str = str(info["error"]) if info["error"] is not None else None
+        return cls(
+            function_id=info["function_id"],
+            function_name=info["function_name"],
+            output=output_str,
+            error=error_str,
+            child_workflow_id=info["child_workflow_id"],
+        )
+
+
 @dataclass
 class ListWorkflowsRequest(BaseMessage):
     body: ListWorkflowsBody
@@ -224,3 +253,14 @@ class ExistPendingWorkflowsRequest(BaseMessage):
 class ExistPendingWorkflowsResponse(BaseMessage):
     exist: bool
     error_message: Optional[str] = None
+
+
+@dataclass
+class ListStepsRequest(BaseMessage):
+    workflow_id: str
+
+
+@dataclass
+class ListStepsResponse(BaseMessage):
+    output: Optional[List[WorkflowSteps]]
+    error_message: Optional[str] = None

dbos/_context.py

@@ -5,6 +5,7 @@ import os
 import uuid
 from contextlib import AbstractContextManager
 from contextvars import ContextVar
+from dataclasses import dataclass
 from enum import Enum
 from types import TracebackType
 from typing import List, Literal, Optional, Type, TypedDict
@@ -48,6 +49,23 @@ class TracedAttributes(TypedDict, total=False):
     authenticatedUserAssumedRole: Optional[str]
 
 
+@dataclass
+class StepStatus:
+    """
+    Status of a step execution.
+
+    Attributes:
+        step_id: The unique ID of this step in its workflow.
+        current_attempt: For steps with automatic retries, which attempt number (zero-indexed) is currently executing.
+        max_attempts: For steps with automatic retries, the maximum number of attempts that will be made before the step fails.
+
+    """
+
+    step_id: int
+    current_attempt: Optional[int]
+    max_attempts: Optional[int]
+
+
 class DBOSContext:
     def __init__(self) -> None:
         self.executor_id = GlobalParams.executor_id
@@ -73,6 +91,7 @@ class DBOSContext:
         self.authenticated_user: Optional[str] = None
         self.authenticated_roles: Optional[List[str]] = None
         self.assumed_role: Optional[str] = None
+        self.step_status: Optional[StepStatus] = None
 
     def create_child(self) -> DBOSContext:
         rv = DBOSContext()
@@ -92,6 +111,9 @@ class DBOSContext:
         rv.assumed_role = self.assumed_role
         return rv
 
+    def has_parent(self) -> bool:
+        return len(self.parent_workflow_id) > 0
+
     def assign_workflow_id(self) -> str:
         if len(self.id_assigned_for_next_workflow) > 0:
             wfid = self.id_assigned_for_next_workflow
@@ -147,10 +169,12 @@ class DBOSContext:
         attributes: TracedAttributes,
     ) -> None:
         self.curr_step_function_id = fid
+        self.step_status = StepStatus(fid, None, None)
         self._start_span(attributes)
 
     def end_step(self, exc_value: Optional[BaseException]) -> None:
         self.curr_step_function_id = -1
+        self.step_status = None
         self._end_span(exc_value)
 
     def start_transaction(
@@ -429,6 +453,30 @@ class EnterDBOSStep:
         return False  # Did not handle
 
 
+class EnterDBOSStepRetry:
+    def __init__(self, current_attempt: int, max_attempts: int) -> None:
+        self.current_attempt = current_attempt
+        self.max_attempts = max_attempts
+
+    def __enter__(self) -> None:
+        ctx = get_local_dbos_context()
+        if ctx is not None and ctx.step_status is not None:
+            ctx.step_status.current_attempt = self.current_attempt
+            ctx.step_status.max_attempts = self.max_attempts
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        ctx = get_local_dbos_context()
+        if ctx is not None and ctx.step_status is not None:
+            ctx.step_status.current_attempt = None
+            ctx.step_status.max_attempts = None
+        return False  # Did not handle
+
+
 class EnterDBOSTransaction:
     def __init__(self, sqls: Session, attributes: TracedAttributes) -> None:
         self.sqls = sqls