dbos 0.28.0a12__tar.gz → 0.28.0a15__tar.gz

{dbos-0.28.0a12 → dbos-0.28.0a15}/PKG-INFO

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

{dbos-0.28.0a12 → dbos-0.28.0a15}/dbos/_app_db.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, TypedDict
+from typing import Any, Dict, List, Optional, TypedDict
 
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
@@ -7,7 +7,6 @@ from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
 from . import _serialization
-from ._dbos_config import DatabaseConfig
 from ._error import DBOSUnexpectedStepError, DBOSWorkflowConflictIDError
 from ._schemas.application_database import ApplicationSchema
 from ._sys_db import StepInfo
@@ -31,67 +30,33 @@ class RecordedResult(TypedDict):
 
 class ApplicationDatabase:
 
-    def __init__(self, database: DatabaseConfig, *, debug_mode: bool = False):
-
-        app_db_name = database["app_db_name"]
+    def __init__(
+        self,
+        *,
+        database_url: str,
+        engine_kwargs: Dict[str, Any],
+        debug_mode: bool = False,
+    ):
+        app_db_url = sa.make_url(database_url).set(drivername="postgresql+psycopg")
 
         # If the application database does not already exist, create it
         if not debug_mode:
-            postgres_db_url = sa.URL.create(
-                "postgresql+psycopg",
-                username=database["username"],
-                password=database["password"],
-                host=database["hostname"],
-                port=database["port"],
-                database="postgres",
+            postgres_db_engine = sa.create_engine(
+                app_db_url.set(database="postgres"),
+                **engine_kwargs,
             )
-            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},
+                    parameters={"db_name": app_db_url.database},
                 ).scalar():
-                    conn.execute(sa.text(f"CREATE DATABASE {app_db_name}"))
+                    conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
             postgres_db_engine.dispose()
 
-        # Create a connection pool for the application database
-        app_db_url = sa.URL.create(
-            "postgresql+psycopg",
-            username=database["username"],
-            password=database["password"],
-            host=database["hostname"],
-            port=database["port"],
-            database=app_db_name,
-        )
-
-        connect_args = {}
-        if (
-            "connectionTimeoutMillis" in database
-            and database["connectionTimeoutMillis"]
-        ):
-            connect_args["connect_timeout"] = int(
-                database["connectionTimeoutMillis"] / 1000
-            )
-
-        pool_size = database.get("app_db_pool_size")
-        if pool_size is None:
-            pool_size = 20
-
-        engine_kwargs = database.get("db_engine_kwargs")
         if engine_kwargs is None:
             engine_kwargs = {}
 
-        # Respect user-provided values. Otherwise, set defaults.
-        if "pool_size" not in engine_kwargs:
-            engine_kwargs["pool_size"] = pool_size
-        if "max_overflow" not in engine_kwargs:
-            engine_kwargs["max_overflow"] = 0
-        if "pool_timeout" not in engine_kwargs:
-            engine_kwargs["pool_timeout"] = 30
-        if "connect_args" not in engine_kwargs:
-            engine_kwargs["connect_args"] = connect_args
-
         self.engine = sa.create_engine(
             app_db_url,
             **engine_kwargs,

{dbos-0.28.0a12 → dbos-0.28.0a15}/dbos/_client.py

@@ -15,7 +15,6 @@ else:
 
 from dbos import _serialization
 from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
-from dbos._dbos_config import parse_database_url_to_dbconfig
 from dbos._error import DBOSException, DBOSNonExistentWorkflowError
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
 from dbos._serialization import WorkflowInputs
@@ -100,11 +99,23 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
 
 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)
-        self._app_db = ApplicationDatabase(db_config)
+        self._sys_db = SystemDatabase(
+            database_url=database_url,
+            engine_kwargs={
+                "pool_timeout": 30,
+                "max_overflow": 0,
+                "pool_size": 2,
+            },
+            sys_db_name=system_database,
+        )
+        self._app_db = ApplicationDatabase(
+            database_url=database_url,
+            engine_kwargs={
+                "pool_timeout": 30,
+                "max_overflow": 0,
+                "pool_size": 2,
+            },
+        )
         self._db_url = database_url
 
     def destroy(self) -> None:

{dbos-0.28.0a12 → dbos-0.28.0a15}/dbos/_conductor/conductor.py

@@ -2,6 +2,7 @@ import socket
 import threading
 import time
 import traceback
+import uuid
 from importlib.metadata import version
 from typing import TYPE_CHECKING, Optional
 
@@ -9,6 +10,7 @@ from websockets import ConnectionClosed, ConnectionClosedOK, InvalidStatus
 from websockets.sync.client import connect
 from websockets.sync.connection import Connection
 
+from dbos._context import SetWorkflowID
 from dbos._utils import GlobalParams
 from dbos._workflow_commands import (
     get_workflow,
@@ -168,10 +170,11 @@ class ConductorWebsocket(threading.Thread):
                             )
                             websocket.send(resume_response.to_json())
                         elif msg_type == p.MessageType.RESTART:
+                            # TODO: deprecate this message type in favor of Fork
                             restart_message = p.RestartRequest.from_json(message)
                             success = True
                             try:
-                                self.dbos.restart_workflow(restart_message.workflow_id)
+                                self.dbos.fork_workflow(restart_message.workflow_id, 1)
                             except Exception as e:
                                 error_message = f"Exception encountered when restarting workflow {restart_message.workflow_id}: {traceback.format_exc()}"
                                 self.dbos.logger.error(error_message)
@@ -183,6 +186,34 @@ class ConductorWebsocket(threading.Thread):
                                 error_message=error_message,
                             )
                             websocket.send(restart_response.to_json())
+                        elif msg_type == p.MessageType.FORK_WORKFLOW:
+                            fork_message = p.ForkWorkflowRequest.from_json(message)
+                            new_workflow_id = fork_message.body["new_workflow_id"]
+                            if new_workflow_id is None:
+                                new_workflow_id = str(uuid.uuid4())
+                            workflow_id = fork_message.body["workflow_id"]
+                            start_step = fork_message.body["start_step"]
+                            app_version = fork_message.body["application_version"]
+                            try:
+                                with SetWorkflowID(new_workflow_id):
+                                    new_handle = self.dbos.fork_workflow(
+                                        workflow_id,
+                                        start_step,
+                                        application_version=app_version,
+                                    )
+                                new_workflow_id = new_handle.workflow_id
+                            except Exception as e:
+                                error_message = f"Exception encountered when forking workflow {workflow_id} to new workflow {new_workflow_id} on step {start_step}, app version {app_version}: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+                                new_workflow_id = None
+
+                            fork_response = p.ForkWorkflowResponse(
+                                type=p.MessageType.FORK_WORKFLOW,
+                                request_id=base_message.request_id,
+                                new_workflow_id=new_workflow_id,
+                                error_message=error_message,
+                            )
+                            websocket.send(fork_response.to_json())
                         elif msg_type == p.MessageType.LIST_WORKFLOWS:
                             list_workflows_message = p.ListWorkflowsRequest.from_json(
                                 message

{dbos-0.28.0a12 → dbos-0.28.0a15}/dbos/_conductor/protocol.py

@@ -17,6 +17,7 @@ class MessageType(str, Enum):
     GET_WORKFLOW = "get_workflow"
     EXIST_PENDING_WORKFLOWS = "exist_pending_workflows"
     LIST_STEPS = "list_steps"
+    FORK_WORKFLOW = "fork_workflow"
 
 
 T = TypeVar("T", bound="BaseMessage")
@@ -133,7 +134,6 @@ class WorkflowsOutput:
     AuthenticatedRoles: Optional[str]
     Input: Optional[str]
     Output: Optional[str]
-    Request: Optional[str]
     Error: Optional[str]
     CreatedAt: Optional[str]
     UpdatedAt: Optional[str]
@@ -167,7 +167,6 @@ class WorkflowsOutput:
             AuthenticatedRoles=roles_str,
             Input=inputs_str,
             Output=outputs_str,
-            Request=request_str,
             Error=error_str,
             CreatedAt=created_at_str,
             UpdatedAt=updated_at_str,
@@ -263,3 +262,21 @@ class ListStepsRequest(BaseMessage):
 class ListStepsResponse(BaseMessage):
     output: Optional[List[WorkflowSteps]]
     error_message: Optional[str] = None
+
+
+class ForkWorkflowBody(TypedDict):
+    workflow_id: str
+    start_step: int
+    application_version: Optional[str]
+    new_workflow_id: Optional[str]
+
+
+@dataclass
+class ForkWorkflowRequest(BaseMessage):
+    body: ForkWorkflowBody
+
+
+@dataclass
+class ForkWorkflowResponse(BaseMessage):
+    new_workflow_id: Optional[str]
+    error_message: Optional[str] = None

{dbos-0.28.0a12 → dbos-0.28.0a15}/dbos/_dbos.py

@@ -24,8 +24,6 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
-    Union,
-    cast,
 )
 
 from opentelemetry.trace import Span
@@ -64,7 +62,6 @@ from ._registrations import (
 )
 from ._roles import default_required_roles, required_roles
 from ._scheduler import ScheduledWorkflow, scheduled
-from ._schemas.system_database import SystemSchema
 from ._sys_db import StepInfo, SystemDatabase, WorkflowStatus, reset_system_database
 from ._tracer import DBOSTracer, dbos_tracer
 
@@ -73,7 +70,7 @@ if TYPE_CHECKING:
     from ._kafka import _KafkaConsumerWorkflow
     from flask import Flask
 
-from sqlalchemy import URL
+from sqlalchemy import make_url
 from sqlalchemy.orm import Session
 
 if sys.version_info < (3, 10):
@@ -418,11 +415,19 @@ class DBOS:
             dbos_logger.info(f"Application version: {GlobalParams.app_version}")
             self._executor_field = ThreadPoolExecutor(max_workers=64)
             self._background_event_loop.start()
+            assert self._config["database_url"] is not None
+            assert self._config["database"]["sys_db_engine_kwargs"] is not None
             self._sys_db_field = SystemDatabase(
-                self._config["database"], debug_mode=debug_mode
+                database_url=self._config["database_url"],
+                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
             self._app_db_field = ApplicationDatabase(
-                self._config["database"], debug_mode=debug_mode
+                database_url=self._config["database_url"],
+                engine_kwargs=self._config["database"]["db_engine_kwargs"],
+                debug_mode=debug_mode,
             )
 
             if debug_mode:
@@ -530,21 +535,13 @@ class DBOS:
             not self._launched
         ), "The system database cannot be reset after DBOS is launched. Resetting the system database is a destructive operation that should only be used in a test environment."
 
-        sysdb_name = (
-            self._config["database"]["sys_db_name"]
-            if "sys_db_name" in self._config["database"]
-            and self._config["database"]["sys_db_name"]
-            else self._config["database"]["app_db_name"] + SystemSchema.sysdb_suffix
-        )
-        postgres_db_url = URL.create(
-            "postgresql+psycopg",
-            username=self._config["database"]["username"],
-            password=self._config["database"]["password"],
-            host=self._config["database"]["hostname"],
-            port=self._config["database"]["port"],
-            database="postgres",
-        )
-        reset_system_database(postgres_db_url, sysdb_name)
+        sysdb_name = self._config["database"]["sys_db_name"]
+        assert sysdb_name is not None
+
+        assert self._config["database_url"] is not None
+        pg_db_url = make_url(self._config["database_url"]).set(database="postgres")
+
+        reset_system_database(pg_db_url, sysdb_name)
 
     def _destroy(self) -> None:
         self._initialized = False