dbos 0.25.1__tar.gz → 0.26.0__tar.gz

{dbos-0.25.1 → dbos-0.26.0}/PKG-INFO

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

{dbos-0.25.1 → dbos-0.26.0}/dbos/__init__.py

@@ -5,8 +5,7 @@ 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
+from ._sys_db import GetWorkflowsInput, WorkflowStatus, WorkflowStatusString
 
 __all__ = [
     "ConfigFile",

{dbos-0.25.1 → dbos-0.26.0}/dbos/_admin_server.py

@@ -7,8 +7,10 @@ from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import TYPE_CHECKING, Any, List, TypedDict
 
+from ._error import DBOSException
 from ._logger import dbos_logger
 from ._recovery import recover_pending_workflows
+from ._utils import GlobalParams
 
 if TYPE_CHECKING:
     from ._dbos import DBOS
@@ -21,6 +23,7 @@ _workflow_queues_metadata_path = "/dbos-workflow-queues-metadata"
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
 # /workflows/:workflow_id/steps
+# /workflows/:workflow_id/fork
 
 
 class AdminServer:
@@ -44,6 +47,7 @@ class AdminServer:
 class AdminRequestHandler(BaseHTTPRequestHandler):
     def __init__(self, dbos: DBOS, *args: Any, **kwargs: Any) -> None:
         self.dbos = dbos
+        self.is_deactivated = False
         super().__init__(*args, **kwargs)
 
     def _end_headers(self) -> None:
@@ -59,10 +63,14 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self._end_headers()
             self.wfile.write("healthy".encode("utf-8"))
         elif self.path == _deactivate_path:
+            if not self.is_deactivated:
+                dbos_logger.info(
+                    f"Deactivating DBOS executor {GlobalParams.executor_id} with version {GlobalParams.app_version}. This executor will complete existing workflows but will not start new workflows."
+                )
+                self.is_deactivated = True
             # Stop all scheduled workflows, queues, and kafka loops
             for event in self.dbos.stop_events:
                 event.set()
-
             self.send_response(200)
             self._end_headers()
             self.wfile.write("deactivated".encode("utf-8"))
@@ -117,6 +125,9 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             restart_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/restart$", self.path
             )
+            fork_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/fork$", self.path
+            )
             resume_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/resume$", self.path
             )
@@ -124,7 +135,23 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
                 r"^/workflows/(?P<workflow_id>[^/]+)/cancel$", self.path
             )
 
-            if restart_match:
+            if fork_match:
+                workflow_id = fork_match.group("workflow_id")
+                try:
+                    data = json.loads(post_data.decode("utf-8"))
+                    start_step: int = data.get("start_step", 1)
+                    self._handle_fork(workflow_id, start_step)
+                except (json.JSONDecodeError, AttributeError) as e:
+                    self.send_response(500)
+                    self.send_header("Content-Type", "application/json")
+                    self.end_headers()
+                    self.wfile.write(
+                        json.dumps({"error": f"Invalid JSON input: {str(e)}"}).encode(
+                            "utf-8"
+                        )
+                    )
+                    return
+            elif restart_match:
                 workflow_id = restart_match.group("workflow_id")
                 self._handle_restart(workflow_id)
             elif resume_match:
@@ -141,10 +168,33 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         return  # Disable admin server request logging
 
     def _handle_restart(self, workflow_id: str) -> None:
-        self.dbos.restart_workflow(workflow_id)
-        print("Restarting workflow", workflow_id)
-        self.send_response(204)
-        self._end_headers()
+        try:
+            print(f"Restarting workflow {workflow_id}")
+            self.dbos.restart_workflow(workflow_id)
+            self.send_response(204)
+            self._end_headers()
+        except DBOSException as e:
+            print(f"Error restarting workflow: {e}")
+            self.send_response(500)
+            response_body = json.dumps({"error": str(e)}).encode("utf-8")
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response_body)))
+            self.end_headers()
+            self.wfile.write(response_body)
+
+    def _handle_fork(self, workflow_id: str, start_step: int) -> None:
+        try:
+            self.dbos.fork_workflow(workflow_id, start_step)
+            self.send_response(204)
+            self._end_headers()
+        except DBOSException as e:
+            print(f"Error forking workflow: {e}")
+            self.send_response(500)
+            response_body = json.dumps({"error": str(e)}).encode("utf-8")
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(response_body)))
+            self.end_headers()
+            self.wfile.write(response_body)
 
     def _handle_resume(self, workflow_id: str) -> None:
         print("Resuming workflow", workflow_id)

{dbos-0.25.1 → dbos-0.26.0}/dbos/_app_db.py

@@ -1,13 +1,16 @@
-from typing import Optional, TypedDict
+from typing import List, Optional, TypedDict
 
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
+from sqlalchemy import inspect, text
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
+from . import _serialization
 from ._dbos_config import ConfigFile, DatabaseConfig
-from ._error import DBOSWorkflowConflictIDError
+from ._error import DBOSUnexpectedStepError, DBOSWorkflowConflictIDError
 from ._schemas.application_database import ApplicationSchema
+from ._sys_db import StepInfo
 
 
 class TransactionResultInternal(TypedDict):
@@ -18,6 +21,7 @@ class TransactionResultInternal(TypedDict):
     txn_id: Optional[str]
     txn_snapshot: str
     executor_id: Optional[str]
+    function_name: Optional[str]
 
 
 class RecordedResult(TypedDict):
@@ -70,9 +74,12 @@ class ApplicationDatabase:
                 database["connectionTimeoutMillis"] / 1000
             )
 
+        pool_size = database.get("app_db_pool_size")
+        if pool_size is None:
+            pool_size = 20
         self.engine = sa.create_engine(
             app_db_url,
-            pool_size=database["app_db_pool_size"],
+            pool_size=pool_size,
             max_overflow=0,
             pool_timeout=30,
             connect_args=connect_args,
@@ -87,7 +94,30 @@ class ApplicationDatabase:
                     f"CREATE SCHEMA IF NOT EXISTS {ApplicationSchema.schema}"
                 )
                 conn.execute(schema_creation_query)
-            ApplicationSchema.metadata_obj.create_all(self.engine)
+
+            inspector = inspect(self.engine)
+            if not inspector.has_table(
+                "transaction_outputs", schema=ApplicationSchema.schema
+            ):
+                ApplicationSchema.metadata_obj.create_all(self.engine)
+            else:
+                columns = inspector.get_columns(
+                    "transaction_outputs", schema=ApplicationSchema.schema
+                )
+                column_names = [col["name"] for col in columns]
+
+                if "function_name" not in column_names:
+                    # Column missing, alter table to add it
+                    with self.engine.connect() as conn:
+                        conn.execute(
+                            text(
+                                f"""
+                            ALTER TABLE {ApplicationSchema.schema}.transaction_outputs
+                            ADD COLUMN function_name TEXT NOT NULL DEFAULT '';
+                            """
+                            )
+                        )
+                        conn.commit()
 
     def destroy(self) -> None:
         self.engine.dispose()
@@ -108,6 +138,7 @@ class ApplicationDatabase:
                     executor_id=(
                         output["executor_id"] if output["executor_id"] else None
                     ),
+                    function_name=output["function_name"],
                 )
             )
         except DBAPIError as dbapi_error:
@@ -133,6 +164,7 @@ class ApplicationDatabase:
                         executor_id=(
                             output["executor_id"] if output["executor_id"] else None
                         ),
+                        function_name=output["function_name"],
                     )
                 )
         except DBAPIError as dbapi_error:
@@ -142,21 +174,116 @@ class ApplicationDatabase:
 
     @staticmethod
     def check_transaction_execution(
-        session: Session, workflow_uuid: str, function_id: int
+        session: Session, workflow_id: str, function_id: int, function_name: str
     ) -> Optional[RecordedResult]:
         rows = session.execute(
             sa.select(
                 ApplicationSchema.transaction_outputs.c.output,
                 ApplicationSchema.transaction_outputs.c.error,
+                ApplicationSchema.transaction_outputs.c.function_name,
             ).where(
-                ApplicationSchema.transaction_outputs.c.workflow_uuid == workflow_uuid,
+                ApplicationSchema.transaction_outputs.c.workflow_uuid == workflow_id,
                 ApplicationSchema.transaction_outputs.c.function_id == function_id,
             )
         ).all()
         if len(rows) == 0:
             return None
+        output, error, recorded_function_name = rows[0][0], rows[0][1], rows[0][2]
+        if function_name != recorded_function_name:
+            raise DBOSUnexpectedStepError(
+                workflow_id=workflow_id,
+                step_id=function_id,
+                expected_name=function_name,
+                recorded_name=recorded_function_name,
+            )
         result: RecordedResult = {
-            "output": rows[0][0],
-            "error": rows[0][1],
+            "output": output,
+            "error": error,
         }
         return result
+
+    def get_transactions(self, workflow_uuid: str) -> List[StepInfo]:
+        with self.engine.begin() as conn:
+            rows = conn.execute(
+                sa.select(
+                    ApplicationSchema.transaction_outputs.c.function_id,
+                    ApplicationSchema.transaction_outputs.c.function_name,
+                    ApplicationSchema.transaction_outputs.c.output,
+                    ApplicationSchema.transaction_outputs.c.error,
+                ).where(
+                    ApplicationSchema.transaction_outputs.c.workflow_uuid
+                    == workflow_uuid,
+                )
+            ).all()
+        return [
+            StepInfo(
+                function_id=row[0],
+                function_name=row[1],
+                output=(
+                    _serialization.deserialize(row[2]) if row[2] is not None else row[2]
+                ),
+                error=(
+                    _serialization.deserialize_exception(row[3])
+                    if row[3] is not None
+                    else row[3]
+                ),
+                child_workflow_id=None,
+            )
+            for row in rows
+        ]
+
+    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
+        with self.engine.begin() as conn:
+            max_function_id_row = conn.execute(
+                sa.select(
+                    sa.func.max(ApplicationSchema.transaction_outputs.c.function_id)
+                ).where(
+                    ApplicationSchema.transaction_outputs.c.workflow_uuid
+                    == workflow_uuid
+                )
+            ).fetchone()
+
+            max_function_id = max_function_id_row[0] if max_function_id_row else None
+
+            return max_function_id
+
+    def clone_workflow_transactions(
+        self, src_workflow_id: str, forked_workflow_id: str, start_step: int
+    ) -> None:
+        """
+        Copies all steps from dbos.transctions_outputs where function_id < input function_id
+        into a new workflow_uuid. Returns the new workflow_uuid.
+        """
+
+        with self.engine.begin() as conn:
+
+            insert_stmt = sa.insert(ApplicationSchema.transaction_outputs).from_select(
+                [
+                    "workflow_uuid",
+                    "function_id",
+                    "output",
+                    "error",
+                    "txn_id",
+                    "txn_snapshot",
+                    "executor_id",
+                    "function_name",
+                ],
+                sa.select(
+                    sa.literal(forked_workflow_id).label("workflow_uuid"),
+                    ApplicationSchema.transaction_outputs.c.function_id,
+                    ApplicationSchema.transaction_outputs.c.output,
+                    ApplicationSchema.transaction_outputs.c.error,
+                    ApplicationSchema.transaction_outputs.c.txn_id,
+                    ApplicationSchema.transaction_outputs.c.txn_snapshot,
+                    ApplicationSchema.transaction_outputs.c.executor_id,
+                    ApplicationSchema.transaction_outputs.c.function_name,
+                ).where(
+                    (
+                        ApplicationSchema.transaction_outputs.c.workflow_uuid
+                        == src_workflow_id
+                    )
+                    & (ApplicationSchema.transaction_outputs.c.function_id < start_step)
+                ),
+            )
+
+            conn.execute(insert_stmt)

{dbos-0.25.1 → dbos-0.26.0}/dbos/_client.py

@@ -1,7 +1,10 @@
 import asyncio
 import sys
+import time
 import uuid
-from typing import Any, Generic, Optional, TypedDict, TypeVar
+from typing import Any, Generic, List, Optional, TypedDict, TypeVar
+
+from dbos._app_db import ApplicationDatabase
 
 if sys.version_info < (3, 11):
     from typing_extensions import NotRequired
@@ -14,18 +17,30 @@ 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
+from dbos._sys_db import (
+    StepInfo,
+    SystemDatabase,
+    WorkflowStatus,
+    WorkflowStatusInternal,
+    WorkflowStatusString,
+)
+from dbos._workflow_commands import (
+    fork_workflow,
+    get_workflow,
+    list_queued_workflows,
+    list_workflow_steps,
+    list_workflows,
+)
 
 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]
+    app_version: NotRequired[str]
+    workflow_timeout: NotRequired[float]
 
 
 class WorkflowHandleClientPolling(Generic[R]):
@@ -41,7 +56,7 @@ class WorkflowHandleClientPolling(Generic[R]):
         res: R = self._sys_db.await_workflow_result(self.workflow_id)
         return res
 
-    def get_status(self) -> "WorkflowStatus":
+    def get_status(self) -> WorkflowStatus:
         status = get_workflow(self._sys_db, self.workflow_id, True)
         if status is None:
             raise DBOSNonExistentWorkflowError(self.workflow_id)
@@ -63,7 +78,7 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
         )
         return res
 
-    async def get_status(self) -> "WorkflowStatus":
+    async def get_status(self) -> WorkflowStatus:
         status = await asyncio.to_thread(
             get_workflow, self._sys_db, self.workflow_id, True
         )
@@ -78,6 +93,7 @@ class DBOSClient:
         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)
 
     def destroy(self) -> None:
         self._sys_db.destroy()
@@ -86,7 +102,6 @@ class DBOSClient:
         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:
@@ -94,12 +109,13 @@ class DBOSClient:
         workflow_id = options.get("workflow_id")
         if workflow_id is None:
             workflow_id = str(uuid.uuid4())
+        workflow_timeout = options.get("workflow_timeout", None)
 
         status: WorkflowStatusInternal = {
             "workflow_uuid": workflow_id,
             "status": WorkflowStatusString.ENQUEUED.value,
             "name": workflow_name,
-            "class_name": workflow_class_name,
+            "class_name": None,
             "queue_name": queue_name,
             "app_version": app_version,
             "config_name": None,
@@ -114,6 +130,10 @@ class DBOSClient:
             "executor_id": None,
             "recovery_attempts": None,
             "app_id": None,
+            "workflow_timeout_ms": (
+                int(workflow_timeout * 1000) if workflow_timeout is not None else None
+            ),
+            "workflow_deadline_epoch_ms": None,
         }
 
         inputs: WorkflowInputs = {
@@ -121,12 +141,9 @@ class DBOSClient:
             "kwargs": kwargs,
         }
 
-        wf_status = self._sys_db.insert_workflow_status(status)
-        self._sys_db.update_workflow_inputs(
-            workflow_id, _serialization.serialize_args(inputs)
+        self._sys_db.init_workflow(
+            status, _serialization.serialize_args(inputs), max_recovery_attempts=None
         )
-        if wf_status == WorkflowStatusString.ENQUEUED.value:
-            self._sys_db.enqueue(workflow_id, queue_name)
         return workflow_id
 
     def enqueue(
@@ -180,8 +197,13 @@ class DBOSClient:
             "recovery_attempts": None,
             "app_id": None,
             "app_version": None,
+            "workflow_timeout_ms": None,
+            "workflow_deadline_epoch_ms": None,
         }
-        self._sys_db.insert_workflow_status(status)
+        with self._sys_db.engine.begin() as conn:
+            self._sys_db.insert_workflow_status(
+                status, conn, max_recovery_attempts=None
+            )
         self._sys_db.send(status["workflow_uuid"], 0, destination_id, message, topic)
 
     async def send_async(
@@ -204,3 +226,141 @@ class DBOSClient:
         return await asyncio.to_thread(
             self.get_event, workflow_id, key, timeout_seconds
         )
+
+    def cancel_workflow(self, workflow_id: str) -> None:
+        self._sys_db.cancel_workflow(workflow_id)
+
+    async def cancel_workflow_async(self, workflow_id: str) -> None:
+        await asyncio.to_thread(self.cancel_workflow, workflow_id)
+
+    def resume_workflow(self, workflow_id: str) -> None:
+        self._sys_db.resume_workflow(workflow_id)
+
+    async def resume_workflow_async(self, workflow_id: str) -> None:
+        await asyncio.to_thread(self.resume_workflow, workflow_id)
+
+    def list_workflows(
+        self,
+        *,
+        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,
+        workflow_id_prefix: Optional[str] = None,
+    ) -> List[WorkflowStatus]:
+        return list_workflows(
+            self._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,
+            workflow_id_prefix=workflow_id_prefix,
+        )
+
+    async def list_workflows_async(
+        self,
+        *,
+        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 await asyncio.to_thread(
+            self.list_workflows,
+            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,
+        )
+
+    def list_queued_workflows(
+        self,
+        *,
+        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(
+            self._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,
+        )
+
+    async def list_queued_workflows_async(
+        self,
+        *,
+        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 await asyncio.to_thread(
+            self.list_queued_workflows,
+            queue_name=queue_name,
+            status=status,
+            start_time=start_time,
+            end_time=end_time,
+            name=name,
+            limit=limit,
+            offset=offset,
+            sort_desc=sort_desc,
+        )
+
+    def list_workflow_steps(self, workflow_id: str) -> List[StepInfo]:
+        return list_workflow_steps(self._sys_db, self._app_db, workflow_id)
+
+    async def list_workflow_steps_async(self, workflow_id: str) -> List[StepInfo]:
+        return await asyncio.to_thread(self.list_workflow_steps, workflow_id)
+
+    def fork_workflow(self, workflow_id: str, start_step: int) -> WorkflowHandle[R]:
+        forked_workflow_id = fork_workflow(
+            self._sys_db, self._app_db, workflow_id, start_step
+        )
+        return WorkflowHandleClientPolling[R](forked_workflow_id, self._sys_db)
+
+    async def fork_workflow_async(
+        self, workflow_id: str, start_step: int
+    ) -> WorkflowHandleAsync[R]:
+        forked_workflow_id = await asyncio.to_thread(
+            fork_workflow, self._sys_db, self._app_db, workflow_id, start_step
+        )
+        return WorkflowHandleClientAsyncPolling[R](forked_workflow_id, self._sys_db)

{dbos-0.25.1 → dbos-0.26.0}/dbos/_conductor/conductor.py

@@ -67,7 +67,7 @@ class ConductorWebsocket(threading.Thread):
                             recovery_message = p.RecoveryRequest.from_json(message)
                             success = True
                             try:
-                                self.dbos.recover_pending_workflows(
+                                self.dbos._recover_pending_workflows(
                                     recovery_message.executor_ids
                                 )
                             except Exception as e:
@@ -254,6 +254,7 @@ class ConductorWebsocket(threading.Thread):
                             try:
                                 step_info = list_workflow_steps(
                                     self.dbos._sys_db,
+                                    self.dbos._app_db,
                                     list_steps_message.workflow_id,
                                 )
                             except Exception as e:

{dbos-0.25.1 → dbos-0.26.0}/dbos/_conductor/protocol.py

@@ -3,8 +3,7 @@ from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import List, Optional, Type, TypedDict, TypeVar
 
-from dbos._sys_db import StepInfo
-from dbos._workflow_commands import WorkflowStatus
+from dbos._sys_db import StepInfo, WorkflowStatus
 
 
 class MessageType(str, Enum):