dbos 0.23.0a14__py3-none-any.whl → 0.24.0__py3-none-any.whl

dbos/__init__.py

@@ -1,13 +1,14 @@
 from . import _error as error
 from ._context import DBOSContextEnsure, DBOSContextSetAuth, SetWorkflowID
 from ._dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
-from ._dbos_config import ConfigFile, get_dbos_database_url, load_config
+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
 
 __all__ = [
     "ConfigFile",
+    "DBOSConfig",
     "DBOS",
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",

dbos/_app_db.py

@@ -5,7 +5,7 @@ import sqlalchemy.dialects.postgresql as pg
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
-from ._dbos_config import ConfigFile
+from ._dbos_config import ConfigFile, DatabaseConfig
 from ._error import DBOSWorkflowConflictIDError
 from ._schemas.application_database import ApplicationSchema
 
@@ -61,8 +61,22 @@ class ApplicationDatabase:
             port=config["database"]["port"],
             database=app_db_name,
         )
+
+        connect_args = {}
+        if (
+            "connectionTimeoutMillis" in config["database"]
+            and config["database"]["connectionTimeoutMillis"]
+        ):
+            connect_args["connect_timeout"] = int(
+                config["database"]["connectionTimeoutMillis"] / 1000
+            )
+
         self.engine = sa.create_engine(
-            app_db_url, pool_size=20, max_overflow=5, pool_timeout=30
+            app_db_url,
+            pool_size=config["database"]["app_db_pool_size"],
+            max_overflow=0,
+            pool_timeout=30,
+            connect_args=connect_args,
         )
         self.sessionmaker = sessionmaker(bind=self.engine)
         self.debug_mode = debug_mode

dbos/_conductor/conductor.py

@@ -1,9 +1,10 @@
+import socket
 import threading
 import time
 import traceback
 from typing import TYPE_CHECKING, Optional
 
-from websockets import ConnectionClosed, ConnectionClosedOK
+from websockets import ConnectionClosed, ConnectionClosedOK, InvalidStatus
 from websockets.sync.client import connect
 from websockets.sync.connection import Connection
 
@@ -33,27 +34,31 @@ class ConductorWebsocket(threading.Thread):
     def run(self) -> None:
         while not self.evt.is_set():
             try:
-                with connect(self.url) as websocket:
+                with connect(
+                    self.url, open_timeout=5, logger=self.dbos.logger
+                ) as websocket:
                     self.websocket = websocket
                     while not self.evt.is_set():
                         message = websocket.recv()
                         if not isinstance(message, str):
                             self.dbos.logger.warning(
-                                "Receieved unexpected non-str message"
+                                "Received unexpected non-str message"
                             )
                             continue
                         base_message = p.BaseMessage.from_json(message)
-                        type = base_message.type
-                        if type == p.MessageType.EXECUTOR_INFO:
+                        msg_type = base_message.type
+                        error_message = None
+                        if msg_type == p.MessageType.EXECUTOR_INFO:
                             info_response = p.ExecutorInfoResponse(
                                 type=p.MessageType.EXECUTOR_INFO,
                                 request_id=base_message.request_id,
                                 executor_id=GlobalParams.executor_id,
                                 application_version=GlobalParams.app_version,
+                                hostname=socket.gethostname(),
                             )
                             websocket.send(info_response.to_json())
                             self.dbos.logger.info("Connected to DBOS conductor")
-                        elif type == p.MessageType.RECOVERY:
+                        elif msg_type == p.MessageType.RECOVERY:
                             recovery_message = p.RecoveryRequest.from_json(message)
                             success = True
                             try:
@@ -61,83 +66,89 @@ class ConductorWebsocket(threading.Thread):
                                     recovery_message.executor_ids
                                 )
                             except Exception as e:
-                                self.dbos.logger.error(
-                                    f"Exception encountered when recovering workflows: {traceback.format_exc()}"
-                                )
+                                error_message = f"Exception encountered when recovering workflows: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
                                 success = False
                             recovery_response = p.RecoveryResponse(
                                 type=p.MessageType.RECOVERY,
                                 request_id=base_message.request_id,
                                 success=success,
+                                error_message=error_message,
                             )
                             websocket.send(recovery_response.to_json())
-                        elif type == p.MessageType.CANCEL:
+                        elif msg_type == p.MessageType.CANCEL:
                             cancel_message = p.CancelRequest.from_json(message)
                             success = True
                             try:
                                 self.dbos.cancel_workflow(cancel_message.workflow_id)
                             except Exception as e:
-                                self.dbos.logger.error(
-                                    f"Exception encountered when cancelling workflow {cancel_message.workflow_id}: {traceback.format_exc()}"
-                                )
+                                error_message = f"Exception encountered when cancelling workflow {cancel_message.workflow_id}: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
                                 success = False
                             cancel_response = p.CancelResponse(
                                 type=p.MessageType.CANCEL,
                                 request_id=base_message.request_id,
                                 success=success,
+                                error_message=error_message,
                             )
                             websocket.send(cancel_response.to_json())
-                        elif type == p.MessageType.RESUME:
+                        elif msg_type == p.MessageType.RESUME:
                             resume_message = p.ResumeRequest.from_json(message)
                             success = True
                             try:
                                 self.dbos.resume_workflow(resume_message.workflow_id)
                             except Exception as e:
-                                self.dbos.logger.error(
-                                    f"Exception encountered when resuming workflow {resume_message.workflow_id}: {traceback.format_exc()}"
-                                )
+                                error_message = f"Exception encountered when resuming workflow {resume_message.workflow_id}: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
                                 success = False
                             resume_response = p.ResumeResponse(
                                 type=p.MessageType.RESUME,
                                 request_id=base_message.request_id,
                                 success=success,
+                                error_message=error_message,
                             )
                             websocket.send(resume_response.to_json())
-                        elif type == p.MessageType.RESTART:
+                        elif msg_type == p.MessageType.RESTART:
                             restart_message = p.RestartRequest.from_json(message)
                             success = True
                             try:
                                 self.dbos.restart_workflow(restart_message.workflow_id)
                             except Exception as e:
-                                self.dbos.logger.error(
-                                    f"Exception encountered when restarting workflow {restart_message.workflow_id}: {traceback.format_exc()}"
-                                )
+                                error_message = f"Exception encountered when restarting workflow {restart_message.workflow_id}: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
                                 success = False
                             restart_response = p.RestartResponse(
                                 type=p.MessageType.RESTART,
                                 request_id=base_message.request_id,
                                 success=success,
+                                error_message=error_message,
                             )
                             websocket.send(restart_response.to_json())
-                        elif type == p.MessageType.LIST_WORKFLOWS:
+                        elif msg_type == p.MessageType.LIST_WORKFLOWS:
                             list_workflows_message = p.ListWorkflowsRequest.from_json(
                                 message
                             )
                             body = list_workflows_message.body
-                            infos = list_workflows(
-                                self.dbos._sys_db,
-                                workflow_ids=body["workflow_uuids"],
-                                user=body["authenticated_user"],
-                                start_time=body["start_time"],
-                                end_time=body["end_time"],
-                                status=body["status"],
-                                request=False,
-                                app_version=body["application_version"],
-                                name=body["workflow_name"],
-                                limit=body["limit"],
-                                offset=body["offset"],
-                                sort_desc=body["sort_desc"],
-                            )
+                            infos = []
+                            try:
+                                infos = list_workflows(
+                                    self.dbos._sys_db,
+                                    workflow_ids=body["workflow_uuids"],
+                                    user=body["authenticated_user"],
+                                    start_time=body["start_time"],
+                                    end_time=body["end_time"],
+                                    status=body["status"],
+                                    request=False,
+                                    app_version=body["application_version"],
+                                    name=body["workflow_name"],
+                                    limit=body["limit"],
+                                    offset=body["offset"],
+                                    sort_desc=body["sort_desc"],
+                                )
+                            except Exception as e:
+                                error_message = f"Exception encountered when listing workflows: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+
                             list_workflows_response = p.ListWorkflowsResponse(
                                 type=p.MessageType.LIST_WORKFLOWS,
                                 request_id=base_message.request_id,
@@ -145,25 +156,32 @@ class ConductorWebsocket(threading.Thread):
                                     p.WorkflowsOutput.from_workflow_information(i)
                                     for i in infos
                                 ],
+                                error_message=error_message,
                             )
                             websocket.send(list_workflows_response.to_json())
-                        elif type == p.MessageType.LIST_QUEUED_WORKFLOWS:
+                        elif msg_type == p.MessageType.LIST_QUEUED_WORKFLOWS:
                             list_queued_workflows_message = (
                                 p.ListQueuedWorkflowsRequest.from_json(message)
                             )
                             q_body = list_queued_workflows_message.body
-                            infos = list_queued_workflows(
-                                self.dbos._sys_db,
-                                start_time=q_body["start_time"],
-                                end_time=q_body["end_time"],
-                                status=q_body["status"],
-                                request=False,
-                                name=q_body["workflow_name"],
-                                limit=q_body["limit"],
-                                offset=q_body["offset"],
-                                queue_name=q_body["queue_name"],
-                                sort_desc=q_body["sort_desc"],
-                            )
+                            infos = []
+                            try:
+                                infos = list_queued_workflows(
+                                    self.dbos._sys_db,
+                                    start_time=q_body["start_time"],
+                                    end_time=q_body["end_time"],
+                                    status=q_body["status"],
+                                    request=False,
+                                    name=q_body["workflow_name"],
+                                    limit=q_body["limit"],
+                                    offset=q_body["offset"],
+                                    queue_name=q_body["queue_name"],
+                                    sort_desc=q_body["sort_desc"],
+                                )
+                            except Exception as e:
+                                error_message = f"Exception encountered when listing queued workflows: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+
                             list_queued_workflows_response = (
                                 p.ListQueuedWorkflowsResponse(
                                     type=p.MessageType.LIST_QUEUED_WORKFLOWS,
@@ -172,18 +190,25 @@ class ConductorWebsocket(threading.Thread):
                                         p.WorkflowsOutput.from_workflow_information(i)
                                         for i in infos
                                     ],
+                                    error_message=error_message,
                                 )
                             )
                             websocket.send(list_queued_workflows_response.to_json())
-                        elif type == p.MessageType.GET_WORKFLOW:
+                        elif msg_type == p.MessageType.GET_WORKFLOW:
                             get_workflow_message = p.GetWorkflowRequest.from_json(
                                 message
                             )
-                            info = get_workflow(
-                                self.dbos._sys_db,
-                                get_workflow_message.workflow_id,
-                                getRequest=False,
-                            )
+                            info = None
+                            try:
+                                info = get_workflow(
+                                    self.dbos._sys_db,
+                                    get_workflow_message.workflow_id,
+                                    getRequest=False,
+                                )
+                            except Exception as e:
+                                error_message = f"Exception encountered when getting workflow {get_workflow_message.workflow_id}: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+
                             get_workflow_response = p.GetWorkflowResponse(
                                 type=p.MessageType.GET_WORKFLOW,
                                 request_id=base_message.request_id,
@@ -192,10 +217,43 @@ class ConductorWebsocket(threading.Thread):
                                     if info is not None
                                     else None
                                 ),
+                                error_message=error_message,
                             )
                             websocket.send(get_workflow_response.to_json())
+                        elif msg_type == p.MessageType.EXIST_PENDING_WORKFLOWS:
+                            exist_pending_workflows_message = (
+                                p.ExistPendingWorkflowsRequest.from_json(message)
+                            )
+                            pending_wfs = []
+                            try:
+                                pending_wfs = self.dbos._sys_db.get_pending_workflows(
+                                    exist_pending_workflows_message.executor_id,
+                                    exist_pending_workflows_message.application_version,
+                                )
+                            except Exception as e:
+                                error_message = f"Exception encountered when checking for pending workflows: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+
+                            exist_pending_workflows_response = (
+                                p.ExistPendingWorkflowsResponse(
+                                    type=p.MessageType.EXIST_PENDING_WORKFLOWS,
+                                    request_id=base_message.request_id,
+                                    exist=len(pending_wfs) > 0,
+                                    error_message=error_message,
+                                )
+                            )
+                            websocket.send(exist_pending_workflows_response.to_json())
                         else:
-                            self.dbos.logger.warning(f"Unexpected message type: {type}")
+                            self.dbos.logger.warning(
+                                f"Unexpected message type: {msg_type}"
+                            )
+                            unknown_message = p.BaseResponse(
+                                request_id=base_message.request_id,
+                                type=msg_type,
+                                error_message="Unknown message type",
+                            )
+                            # Still need to send a response to the conductor
+                            websocket.send(unknown_message.to_json())
             except ConnectionClosedOK:
                 self.dbos.logger.info("Conductor connection terminated")
                 break
@@ -205,6 +263,14 @@ class ConductorWebsocket(threading.Thread):
                 )
                 time.sleep(1)
                 continue
+            except InvalidStatus as e:
+                # This happens when it cannot open a connection to the conductor. E.g., the conductor rejects the request
+                json_data = e.response.body.decode("utf-8")
+                self.dbos.logger.error(
+                    f"Failed to connect to conductor. Retrying: {str(e) }. Details: {json_data}"
+                )
+                time.sleep(1)
+                continue
             except Exception as e:
                 self.dbos.logger.error(
                     f"Unexpected exception in connection to conductor. Reconnecting: {e}"

dbos/_conductor/protocol.py

@@ -15,6 +15,7 @@ class MessageType(str, Enum):
     RESUME = "resume"
     RESTART = "restart"
     GET_WORKFLOW = "get_workflow"
+    EXIST_PENDING_WORKFLOWS = "exist_pending_workflows"
 
 
 T = TypeVar("T", bound="BaseMessage")
@@ -44,6 +45,11 @@ class BaseMessage:
         return json.dumps(dict_data)
 
 
+@dataclass
+class BaseResponse(BaseMessage):
+    error_message: Optional[str] = None
+
+
 @dataclass
 class ExecutorInfoRequest(BaseMessage):
     pass
@@ -53,6 +59,8 @@ class ExecutorInfoRequest(BaseMessage):
 class ExecutorInfoResponse(BaseMessage):
     executor_id: str
     application_version: str
+    hostname: Optional[str]
+    error_message: Optional[str] = None
 
 
 @dataclass
@@ -63,6 +71,7 @@ class RecoveryRequest(BaseMessage):
 @dataclass
 class RecoveryResponse(BaseMessage):
     success: bool
+    error_message: Optional[str] = None
 
 
 @dataclass
@@ -73,6 +82,7 @@ class CancelRequest(BaseMessage):
 @dataclass
 class CancelResponse(BaseMessage):
     success: bool
+    error_message: Optional[str] = None
 
 
 @dataclass
@@ -83,6 +93,7 @@ class ResumeRequest(BaseMessage):
 @dataclass
 class ResumeResponse(BaseMessage):
     success: bool
+    error_message: Optional[str] = None
 
 
 @dataclass
@@ -93,6 +104,7 @@ class RestartRequest(BaseMessage):
 @dataclass
 class RestartResponse(BaseMessage):
     success: bool
+    error_message: Optional[str] = None
 
 
 class ListWorkflowsBody(TypedDict):
@@ -126,6 +138,7 @@ class WorkflowsOutput:
     UpdatedAt: Optional[str]
     QueueName: Optional[str]
     ApplicationVersion: Optional[str]
+    ExecutorID: Optional[str]
 
     @classmethod
     def from_workflow_information(cls, info: WorkflowInformation) -> "WorkflowsOutput":
@@ -153,6 +166,7 @@ class WorkflowsOutput:
             UpdatedAt=updated_at_str,
             QueueName=info.queue_name,
             ApplicationVersion=info.app_version,
+            ExecutorID=info.executor_id,
         )
 
 
@@ -164,6 +178,7 @@ class ListWorkflowsRequest(BaseMessage):
 @dataclass
 class ListWorkflowsResponse(BaseMessage):
     output: List[WorkflowsOutput]
+    error_message: Optional[str] = None
 
 
 class ListQueuedWorkflowsBody(TypedDict):
@@ -185,6 +200,7 @@ class ListQueuedWorkflowsRequest(BaseMessage):
 @dataclass
 class ListQueuedWorkflowsResponse(BaseMessage):
     output: List[WorkflowsOutput]
+    error_message: Optional[str] = None
 
 
 @dataclass
@@ -195,3 +211,16 @@ class GetWorkflowRequest(BaseMessage):
 @dataclass
 class GetWorkflowResponse(BaseMessage):
     output: Optional[WorkflowsOutput]
+    error_message: Optional[str] = None
+
+
+@dataclass
+class ExistPendingWorkflowsRequest(BaseMessage):
+    executor_id: str
+    application_version: str
+
+
+@dataclass
+class ExistPendingWorkflowsResponse(BaseMessage):
+    exist: bool
+    error_message: Optional[str] = None