dbos 0.22.0a11__py3-none-any.whl → 0.23.0__py3-none-any.whl

dbos/__main__.py

@@ -0,0 +1,26 @@
+import re
+import sys
+from typing import NoReturn, Optional, Union
+
+from dbos.cli.cli import app
+
+
+def main() -> NoReturn:
+    # Modify sys.argv[0] to remove script or executable extensions
+    sys.argv[0] = re.sub(r"(-script\.pyw|\.exe)?$", "", sys.argv[0])
+
+    retval: Optional[Union[str, int]] = 1
+    try:
+        app()
+        retval = None
+    except SystemExit as e:
+        retval = e.code
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        retval = 1
+    finally:
+        sys.exit(retval)
+
+
+if __name__ == "__main__":
+    main()

dbos/_app_db.py

@@ -27,29 +27,30 @@ class RecordedResult(TypedDict):
 
 class ApplicationDatabase:
 
-    def __init__(self, config: ConfigFile):
+    def __init__(self, config: ConfigFile, *, debug_mode: bool = False):
         self.config = config
 
         app_db_name = config["database"]["app_db_name"]
 
         # If the application database does not already exist, create it
-        postgres_db_url = sa.URL.create(
-            "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
-            database="postgres",
-        )
-        postgres_db_engine = sa.create_engine(postgres_db_url)
-        with postgres_db_engine.connect() as conn:
-            conn.execution_options(isolation_level="AUTOCOMMIT")
-            if not conn.execute(
-                sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
-                parameters={"db_name": app_db_name},
-            ).scalar():
-                conn.execute(sa.text(f"CREATE DATABASE {app_db_name}"))
-        postgres_db_engine.dispose()
+        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"],
+                database="postgres",
+            )
+            postgres_db_engine = sa.create_engine(postgres_db_url)
+            with postgres_db_engine.connect() as conn:
+                conn.execution_options(isolation_level="AUTOCOMMIT")
+                if not conn.execute(
+                    sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
+                    parameters={"db_name": app_db_name},
+                ).scalar():
+                    conn.execute(sa.text(f"CREATE DATABASE {app_db_name}"))
+            postgres_db_engine.dispose()
 
         # Create a connection pool for the application database
         app_db_url = sa.URL.create(
@@ -64,14 +65,16 @@ class ApplicationDatabase:
             app_db_url, pool_size=20, max_overflow=5, pool_timeout=30
         )
         self.sessionmaker = sessionmaker(bind=self.engine)
+        self.debug_mode = debug_mode
 
         # Create the dbos schema and transaction_outputs table in the application database
-        with self.engine.begin() as conn:
-            schema_creation_query = sa.text(
-                f"CREATE SCHEMA IF NOT EXISTS {ApplicationSchema.schema}"
-            )
-            conn.execute(schema_creation_query)
-        ApplicationSchema.metadata_obj.create_all(self.engine)
+        if not debug_mode:
+            with self.engine.begin() as conn:
+                schema_creation_query = sa.text(
+                    f"CREATE SCHEMA IF NOT EXISTS {ApplicationSchema.schema}"
+                )
+                conn.execute(schema_creation_query)
+            ApplicationSchema.metadata_obj.create_all(self.engine)
 
     def destroy(self) -> None:
         self.engine.dispose()
@@ -100,6 +103,8 @@ class ApplicationDatabase:
             raise
 
     def record_transaction_error(self, output: TransactionResultInternal) -> None:
+        if self.debug_mode:
+            raise Exception("called record_transaction_error in debug mode")
         try:
             with self.engine.begin() as conn:
                 conn.execute(

dbos/_conductor/conductor.py

@@ -0,0 +1,213 @@
+import threading
+import time
+import traceback
+from typing import TYPE_CHECKING, Optional
+
+from websockets import ConnectionClosed, ConnectionClosedOK
+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 . import protocol as p
+
+if TYPE_CHECKING:
+    from dbos import DBOS
+
+
+class ConductorWebsocket(threading.Thread):
+
+    def __init__(
+        self, dbos: "DBOS", conductor_url: str, conductor_key: str, evt: threading.Event
+    ):
+        super().__init__(daemon=True)
+        self.websocket: Optional[Connection] = None
+        self.evt = evt
+        self.dbos = dbos
+        self.app_name = dbos.config["name"]
+        self.url = (
+            conductor_url.rstrip("/") + f"/websocket/{self.app_name}/{conductor_key}"
+        )
+
+    def run(self) -> None:
+        while not self.evt.is_set():
+            try:
+                with connect(self.url) 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"
+                            )
+                            continue
+                        base_message = p.BaseMessage.from_json(message)
+                        type = base_message.type
+                        if 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,
+                            )
+                            websocket.send(info_response.to_json())
+                            self.dbos.logger.info("Connected to DBOS conductor")
+                        elif type == p.MessageType.RECOVERY:
+                            recovery_message = p.RecoveryRequest.from_json(message)
+                            success = True
+                            try:
+                                self.dbos.recover_pending_workflows(
+                                    recovery_message.executor_ids
+                                )
+                            except Exception as e:
+                                self.dbos.logger.error(
+                                    f"Exception encountered when recovering workflows: {traceback.format_exc()}"
+                                )
+                                success = False
+                            recovery_response = p.RecoveryResponse(
+                                type=p.MessageType.RECOVERY,
+                                request_id=base_message.request_id,
+                                success=success,
+                            )
+                            websocket.send(recovery_response.to_json())
+                        elif 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()}"
+                                )
+                                success = False
+                            cancel_response = p.CancelResponse(
+                                type=p.MessageType.CANCEL,
+                                request_id=base_message.request_id,
+                                success=success,
+                            )
+                            websocket.send(cancel_response.to_json())
+                        elif 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()}"
+                                )
+                                success = False
+                            resume_response = p.ResumeResponse(
+                                type=p.MessageType.RESUME,
+                                request_id=base_message.request_id,
+                                success=success,
+                            )
+                            websocket.send(resume_response.to_json())
+                        elif 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()}"
+                                )
+                                success = False
+                            restart_response = p.RestartResponse(
+                                type=p.MessageType.RESTART,
+                                request_id=base_message.request_id,
+                                success=success,
+                            )
+                            websocket.send(restart_response.to_json())
+                        elif 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"],
+                            )
+                            list_workflows_response = p.ListWorkflowsResponse(
+                                type=p.MessageType.LIST_WORKFLOWS,
+                                request_id=base_message.request_id,
+                                output=[
+                                    p.WorkflowsOutput.from_workflow_information(i)
+                                    for i in infos
+                                ],
+                            )
+                            websocket.send(list_workflows_response.to_json())
+                        elif 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"],
+                            )
+                            list_queued_workflows_response = (
+                                p.ListQueuedWorkflowsResponse(
+                                    type=p.MessageType.LIST_QUEUED_WORKFLOWS,
+                                    request_id=base_message.request_id,
+                                    output=[
+                                        p.WorkflowsOutput.from_workflow_information(i)
+                                        for i in infos
+                                    ],
+                                )
+                            )
+                            websocket.send(list_queued_workflows_response.to_json())
+                        elif 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,
+                            )
+                            get_workflow_response = p.GetWorkflowResponse(
+                                type=p.MessageType.GET_WORKFLOW,
+                                request_id=base_message.request_id,
+                                output=(
+                                    p.WorkflowsOutput.from_workflow_information(info)
+                                    if info is not None
+                                    else None
+                                ),
+                            )
+                            websocket.send(get_workflow_response.to_json())
+                        else:
+                            self.dbos.logger.warning(f"Unexpected message type: {type}")
+            except ConnectionClosedOK:
+                self.dbos.logger.info("Conductor connection terminated")
+                break
+            except ConnectionClosed as e:
+                self.dbos.logger.warning(
+                    f"Connection to conductor lost. Reconnecting: {e}"
+                )
+                time.sleep(1)
+                continue
+            except Exception as e:
+                self.dbos.logger.error(
+                    f"Unexpected exception in connection to conductor. Reconnecting: {e}"
+                )
+                time.sleep(1)
+                continue

dbos/_conductor/protocol.py

@@ -0,0 +1,197 @@
+import json
+from dataclasses import asdict, dataclass
+from enum import Enum
+from typing import List, Optional, Type, TypedDict, TypeVar
+
+from dbos._workflow_commands import WorkflowInformation
+
+
+class MessageType(str, Enum):
+    EXECUTOR_INFO = "executor_info"
+    RECOVERY = "recovery"
+    CANCEL = "cancel"
+    LIST_WORKFLOWS = "list_workflows"
+    LIST_QUEUED_WORKFLOWS = "list_queued_workflows"
+    RESUME = "resume"
+    RESTART = "restart"
+    GET_WORKFLOW = "get_workflow"
+
+
+T = TypeVar("T", bound="BaseMessage")
+
+
+@dataclass
+class BaseMessage:
+    type: MessageType
+    request_id: str
+
+    @classmethod
+    def from_json(cls: Type[T], json_str: str) -> T:
+        """
+        Safely load a JSON into a dataclass, loading only the
+        attributes specified in the dataclass.
+        """
+        data = json.loads(json_str)
+        all_annotations = {}
+        for base_cls in cls.__mro__:
+            if hasattr(base_cls, "__annotations__"):
+                all_annotations.update(base_cls.__annotations__)
+        kwargs = {k: v for k, v in data.items() if k in all_annotations}
+        return cls(**kwargs)
+
+    def to_json(self) -> str:
+        dict_data = asdict(self)
+        return json.dumps(dict_data)
+
+
+@dataclass
+class ExecutorInfoRequest(BaseMessage):
+    pass
+
+
+@dataclass
+class ExecutorInfoResponse(BaseMessage):
+    executor_id: str
+    application_version: str
+
+
+@dataclass
+class RecoveryRequest(BaseMessage):
+    executor_ids: List[str]
+
+
+@dataclass
+class RecoveryResponse(BaseMessage):
+    success: bool
+
+
+@dataclass
+class CancelRequest(BaseMessage):
+    workflow_id: str
+
+
+@dataclass
+class CancelResponse(BaseMessage):
+    success: bool
+
+
+@dataclass
+class ResumeRequest(BaseMessage):
+    workflow_id: str
+
+
+@dataclass
+class ResumeResponse(BaseMessage):
+    success: bool
+
+
+@dataclass
+class RestartRequest(BaseMessage):
+    workflow_id: str
+
+
+@dataclass
+class RestartResponse(BaseMessage):
+    success: bool
+
+
+class ListWorkflowsBody(TypedDict):
+    workflow_uuids: List[str]
+    workflow_name: Optional[str]
+    authenticated_user: Optional[str]
+    start_time: Optional[str]
+    end_time: Optional[str]
+    status: Optional[str]
+    application_version: Optional[str]
+    limit: Optional[int]
+    offset: Optional[int]
+    sort_desc: bool
+
+
+@dataclass
+class WorkflowsOutput:
+    WorkflowUUID: str
+    Status: Optional[str]
+    WorkflowName: Optional[str]
+    WorkflowClassName: Optional[str]
+    WorkflowConfigName: Optional[str]
+    AuthenticatedUser: Optional[str]
+    AssumedRole: Optional[str]
+    AuthenticatedRoles: Optional[str]
+    Input: Optional[str]
+    Output: Optional[str]
+    Request: Optional[str]
+    Error: Optional[str]
+    CreatedAt: Optional[str]
+    UpdatedAt: Optional[str]
+    QueueName: Optional[str]
+    ApplicationVersion: Optional[str]
+
+    @classmethod
+    def from_workflow_information(cls, info: WorkflowInformation) -> "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
+        request_str = str(info.request) if info.request 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,
+            AuthenticatedUser=info.authenticated_user,
+            AssumedRole=info.assumed_role,
+            AuthenticatedRoles=info.authenticated_roles,
+            Input=inputs_str,
+            Output=outputs_str,
+            Request=request_str,
+            Error=info.error,
+            CreatedAt=created_at_str,
+            UpdatedAt=updated_at_str,
+            QueueName=info.queue_name,
+            ApplicationVersion=info.app_version,
+        )
+
+
+@dataclass
+class ListWorkflowsRequest(BaseMessage):
+    body: ListWorkflowsBody
+
+
+@dataclass
+class ListWorkflowsResponse(BaseMessage):
+    output: List[WorkflowsOutput]
+
+
+class ListQueuedWorkflowsBody(TypedDict):
+    workflow_name: Optional[str]
+    start_time: Optional[str]
+    end_time: Optional[str]
+    status: Optional[str]
+    queue_name: Optional[str]
+    limit: Optional[int]
+    offset: Optional[int]
+    sort_desc: bool
+
+
+@dataclass
+class ListQueuedWorkflowsRequest(BaseMessage):
+    body: ListQueuedWorkflowsBody
+
+
+@dataclass
+class ListQueuedWorkflowsResponse(BaseMessage):
+    output: List[WorkflowsOutput]
+
+
+@dataclass
+class GetWorkflowRequest(BaseMessage):
+    workflow_id: str
+
+
+@dataclass
+class GetWorkflowResponse(BaseMessage):
+    output: Optional[WorkflowsOutput]

dbos/_context.py

@@ -12,6 +12,8 @@ from typing import List, Literal, Optional, Type, TypedDict
 from opentelemetry.trace import Span, Status, StatusCode
 from sqlalchemy.orm import Session
 
+from dbos._utils import GlobalParams
+
 from ._logger import dbos_logger
 from ._request import Request
 from ._tracer import dbos_tracer
@@ -48,7 +50,7 @@ class TracedAttributes(TypedDict, total=False):
 
 class DBOSContext:
     def __init__(self) -> None:
-        self.executor_id = os.environ.get("DBOS__VMID", "local")
+        self.executor_id = GlobalParams.executor_id
         self.app_id = os.environ.get("DBOS__APPID", "")
 
         self.logger = dbos_logger