dbos 0.19.0a9__tar.gz → 0.20.0a2__tar.gz

{dbos-0.19.0a9 → dbos-0.20.0a2}/PKG-INFO

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

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/_admin_server.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import re
 import threading
 from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -15,6 +16,9 @@ if TYPE_CHECKING:
 _health_check_path = "/dbos-healthz"
 _workflow_recovery_path = "/dbos-workflow-recovery"
 _deactivate_path = "/deactivate"
+# /workflows/:workflow_id/cancel
+# /workflows/:workflow_id/resume
+# /workflows/:workflow_id/restart
 
 
 class AdminServer:
@@ -79,12 +83,51 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self._end_headers()
             self.wfile.write(json.dumps(workflow_ids).encode("utf-8"))
         else:
-            self.send_response(404)
-            self._end_headers()
+
+            restart_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/restart$", self.path
+            )
+            resume_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/resume$", self.path
+            )
+            cancel_match = re.match(
+                r"^/workflows/(?P<workflow_id>[^/]+)/cancel$", self.path
+            )
+
+            if restart_match:
+                workflow_id = restart_match.group("workflow_id")
+                self._handle_restart(workflow_id)
+            elif resume_match:
+                workflow_id = resume_match.group("workflow_id")
+                self._handle_resume(workflow_id)
+            elif cancel_match:
+                workflow_id = cancel_match.group("workflow_id")
+                self._handle_cancel(workflow_id)
+            else:
+                self.send_response(404)
+                self._end_headers()
 
     def log_message(self, format: str, *args: Any) -> None:
         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()
+
+    def _handle_resume(self, workflow_id: str) -> None:
+        print("Resuming workflow", workflow_id)
+        self.dbos.resume_workflow(workflow_id)
+        self.send_response(204)
+        self._end_headers()
+
+    def _handle_cancel(self, workflow_id: str) -> None:
+        print("Cancelling workflow", workflow_id)
+        self.dbos.cancel_workflow(workflow_id)
+        self.send_response(204)
+        self._end_headers()
+
 
 # Be consistent with DBOS-TS response.
 class PerfUtilization(TypedDict):

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/_core.py

@@ -266,7 +266,9 @@ def _execute_workflow_wthread(
                 raise
 
 
-def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]":
+def execute_workflow_by_id(
+    dbos: "DBOS", workflow_id: str, startNew: bool = False
+) -> "WorkflowHandle[Any]":
     status = dbos._sys_db.get_workflow_status(workflow_id)
     if not status:
         raise DBOSRecoveryError(workflow_id, "Workflow status not found")
@@ -293,7 +295,8 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     workflow_id,
                     f"Cannot execute workflow because instance '{iname}' is not registered",
                 )
-            with SetWorkflowID(workflow_id):
+
+            if startNew:
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -303,6 +306,17 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
+            else:
+                with SetWorkflowID(workflow_id):
+                    return start_workflow(
+                        dbos,
+                        wf_func,
+                        status["queue_name"],
+                        True,
+                        dbos._registry.instance_info_map[iname],
+                        *inputs["args"],
+                        **inputs["kwargs"],
+                    )
         elif status["class_name"] is not None:
             class_name = status["class_name"]
             if class_name not in dbos._registry.class_info_map:
@@ -310,7 +324,8 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     workflow_id,
                     f"Cannot execute workflow because class '{class_name}' is not registered",
                 )
-            with SetWorkflowID(workflow_id):
+
+            if startNew:
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -320,8 +335,19 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
+            else:
+                with SetWorkflowID(workflow_id):
+                    return start_workflow(
+                        dbos,
+                        wf_func,
+                        status["queue_name"],
+                        True,
+                        dbos._registry.class_info_map[class_name],
+                        *inputs["args"],
+                        **inputs["kwargs"],
+                    )
         else:
-            with SetWorkflowID(workflow_id):
+            if startNew:
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -330,6 +356,16 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
+            else:
+                with SetWorkflowID(workflow_id):
+                    return start_workflow(
+                        dbos,
+                        wf_func,
+                        status["queue_name"],
+                        True,
+                        *inputs["args"],
+                        **inputs["kwargs"],
+                    )
 
 
 @overload

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/_dbos.py

@@ -56,6 +56,7 @@ from ._registrations import (
 )
 from ._roles import default_required_roles, required_roles
 from ._scheduler import ScheduledWorkflow, scheduled
+from ._sys_db import WorkflowStatusString
 from ._tracer import dbos_tracer
 
 if TYPE_CHECKING:
@@ -231,6 +232,7 @@ class DBOS:
                         f"DBOS configured multiple times with conflicting information"
                     )
                 config = _dbos_global_registry.config
+
             _dbos_global_instance = super().__new__(cls)
             _dbos_global_instance.__init__(fastapi=fastapi, config=config, flask=flask)  # type: ignore
         else:
@@ -767,6 +769,11 @@ class DBOS:
         """Execute a workflow by ID (for recovery)."""
         return execute_workflow_by_id(_get_dbos_instance(), workflow_id)
 
+    @classmethod
+    def restart_workflow(cls, workflow_id: str) -> None:
+        """Execute a workflow by ID (for recovery)."""
+        execute_workflow_by_id(_get_dbos_instance(), workflow_id, True)
+
     @classmethod
     def recover_pending_workflows(
         cls, executor_ids: List[str] = ["local"]
@@ -774,6 +781,18 @@ class DBOS:
         """Find all PENDING workflows and execute them."""
         return recover_pending_workflows(_get_dbos_instance(), executor_ids)
 
+    @classmethod
+    def cancel_workflow(cls, workflow_id: str) -> None:
+        """Cancel a workflow by ID."""
+        _get_dbos_instance()._sys_db.set_workflow_status(
+            workflow_id, WorkflowStatusString.CANCELLED, False
+        )
+
+    @classmethod
+    def resume_workflow(cls, workflow_id: str) -> None:
+        """Resume a workflow by ID."""
+        execute_workflow_by_id(_get_dbos_instance(), workflow_id, False)
+
     @classproperty
     def logger(cls) -> Logger:
         """Return the DBOS `Logger` for the current context."""

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/_fastapi.py

@@ -94,7 +94,11 @@ def setup_fastapi_middleware(app: FastAPI, dbos: DBOS) -> None:
         with EnterDBOSHandler(attributes):
             ctx = assert_current_dbos_context()
             ctx.request = _make_request(request)
-            workflow_id = request.headers.get("dbos-idempotency-key", "")
-            with SetWorkflowID(workflow_id):
+            workflow_id = request.headers.get("dbos-idempotency-key")
+            if workflow_id is not None:
+                # Set the workflow ID for the handler
+                with SetWorkflowID(workflow_id):
+                    response = await call_next(request)
+            else:
                 response = await call_next(request)
         return response

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/_flask.py

@@ -34,8 +34,12 @@ class FlaskMiddleware:
         with EnterDBOSHandler(attributes):
             ctx = assert_current_dbos_context()
             ctx.request = _make_request(request)
-            workflow_id = request.headers.get("dbos-idempotency-key", "")
-            with SetWorkflowID(workflow_id):
+            workflow_id = request.headers.get("dbos-idempotency-key")
+            if workflow_id is not None:
+                # Set the workflow ID for the handler
+                with SetWorkflowID(workflow_id):
+                    response = self.app(environ, start_response)
+            else:
                 response = self.app(environ, start_response)
         return response
 

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/_workflow_commands.py

@@ -1,3 +1,6 @@
+import importlib
+import os
+import sys
 from typing import Any, List, Optional, cast
 
 import typer
@@ -6,6 +9,7 @@ from rich import print
 from dbos import DBOS
 
 from . import _serialization, load_config
+from ._core import execute_workflow_by_id
 from ._dbos_config import ConfigFile, _is_valid_app_name
 from ._sys_db import (
     GetWorkflowsInput,
@@ -123,11 +127,6 @@ def _cancel_workflow(config: ConfigFile, uuid: str) -> None:
             sys_db.destroy()
 
 
-def _reattempt_workflow(uuid: str, startNewWorkflow: bool) -> None:
-    print(f"Reattempt workflow info for {uuid} not implemented")
-    return
-
-
 def _get_workflow_info(
     sys_db: SystemDatabase, workflowUUID: str, getRequest: bool
 ) -> Optional[WorkflowInformation]:

{dbos-0.19.0a9 → dbos-0.20.0a2}/dbos/cli.py

@@ -9,6 +9,7 @@ from os import path
 from typing import Any
 
 import jsonpickle  # type: ignore
+import requests
 import sqlalchemy as sa
 import tomlkit
 import typer
@@ -22,12 +23,7 @@ from . import _serialization, load_config
 from ._app_db import ApplicationDatabase
 from ._dbos_config import _is_valid_app_name
 from ._sys_db import SystemDatabase
-from ._workflow_commands import (
-    _cancel_workflow,
-    _get_workflow,
-    _list_workflows,
-    _reattempt_workflow,
-)
+from ._workflow_commands import _cancel_workflow, _get_workflow, _list_workflows
 
 app = typer.Typer()
 workflow = typer.Typer()
@@ -432,5 +428,49 @@ def cancel(
     print(f"Workflow {uuid} has been cancelled")
 
 
+@workflow.command(help="Resume a workflow that has been cancelled")
+def resume(
+    uuid: Annotated[str, typer.Argument()],
+    host: Annotated[
+        typing.Optional[str],
+        typer.Option("--host", "-h", help="Specify the admin host"),
+    ] = "localhost",
+    port: Annotated[
+        typing.Optional[int],
+        typer.Option("--port", "-p", help="Specify the admin port"),
+    ] = 3001,
+) -> None:
+    response = requests.post(
+        f"http://{host}:{port}/workflows/{uuid}/resume", json=[], timeout=5
+    )
+
+    if response.status_code == 200:
+        print(f"Workflow {uuid} has been resumed")
+    else:
+        print(f"Failed to resume workflow {uuid}. Status code: {response.status_code}")
+
+
+@workflow.command(help="Restart a workflow from the beginning with a new id")
+def restart(
+    uuid: Annotated[str, typer.Argument()],
+    host: Annotated[
+        typing.Optional[str],
+        typer.Option("--host", "-h", help="Specify the admin host"),
+    ] = "localhost",
+    port: Annotated[
+        typing.Optional[int],
+        typer.Option("--port", "-p", help="Specify the admin port"),
+    ] = 3001,
+) -> None:
+    response = requests.post(
+        f"http://{host}:{port}/workflows/{uuid}/restart", json=[], timeout=5
+    )
+
+    if response.status_code == 200:
+        print(f"Workflow {uuid} has been restarted")
+    else:
+        print(f"Failed to resume workflow {uuid}. Status code: {response.status_code}")
+
+
 if __name__ == "__main__":
     app()

{dbos-0.19.0a9 → dbos-0.20.0a2}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.19.0a9"
+version = "0.20.0a2"
 
 [project.license]
 text = "MIT"

{dbos-0.19.0a9 → dbos-0.20.0a2}/tests/test_admin_server.py

@@ -5,7 +5,7 @@ import uuid
 import requests
 
 # Public API
-from dbos import DBOS, SetWorkflowID
+from dbos import DBOS, ConfigFile, SetWorkflowID, _workflow_commands
 
 
 def test_admin_endpoints(dbos: DBOS) -> None:
@@ -148,3 +148,121 @@ runtimeConfig:
         # Clean up after the test
         DBOS.destroy()
         os.remove("dbos-config.yaml")
+
+
+def test_admin_workflow_resume(dbos: DBOS, config: ConfigFile) -> None:
+
+    @DBOS.workflow()
+    def simple_workflow() -> None:
+        print("Executed Simple workflow")
+        return
+
+    # run the workflow
+    simple_workflow()
+    time.sleep(1)
+
+    # get the workflow list
+    output = _workflow_commands._list_workflows(
+        config, 10, None, None, None, None, False, None
+    )
+    assert len(output) == 1, f"Expected list length to be 1, but got {len(output)}"
+
+    assert output[0] != None, "Expected output to be not None"
+
+    wfUuid = output[0].workflowUUID
+
+    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    assert info is not None, "Expected output to be not None"
+
+    assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
+
+    response = requests.post(
+        f"http://localhost:3001/workflows/{wfUuid}/cancel", json=[], timeout=5
+    )
+    assert response.status_code == 204
+
+    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    if info is not None:
+        assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
+    else:
+        assert False, "Expected info to be not None"
+
+    response = requests.post(
+        f"http://localhost:3001/workflows/{wfUuid}/resume", json=[], timeout=5
+    )
+    assert response.status_code == 204
+
+    time.sleep(1)
+
+    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    if info is not None:
+        assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
+    else:
+        assert False, "Expected info to be not None"
+
+
+def test_admin_workflow_restart(dbos: DBOS, config: ConfigFile) -> None:
+
+    @DBOS.workflow()
+    def simple_workflow() -> None:
+        print("Executed Simple workflow")
+        return
+
+    # run the workflow
+    simple_workflow()
+    time.sleep(1)
+
+    # get the workflow list
+    output = _workflow_commands._list_workflows(
+        config, 10, None, None, None, None, False, None
+    )
+    assert len(output) == 1, f"Expected list length to be 1, but got {len(output)}"
+
+    assert output[0] != None, "Expected output to be not None"
+
+    wfUuid = output[0].workflowUUID
+
+    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    assert info is not None, "Expected output to be not None"
+
+    assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
+
+    response = requests.post(
+        f"http://localhost:3001/workflows/{wfUuid}/cancel", json=[], timeout=5
+    )
+    assert response.status_code == 204
+
+    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    if info is not None:
+        assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
+    else:
+        assert False, "Expected info to be not None"
+
+    response = requests.post(
+        f"http://localhost:3001/workflows/{wfUuid}/restart", json=[], timeout=5
+    )
+    assert response.status_code == 204
+
+    time.sleep(1)
+
+    info = _workflow_commands._get_workflow(config, wfUuid, True)
+    if info is not None:
+        assert info.status == "CANCELLED", f"Expected status to be CANCELLED"
+    else:
+        assert False, "Expected info to be not None"
+
+    output = _workflow_commands._list_workflows(
+        config, 10, None, None, None, None, False, None
+    )
+    assert len(output) == 2, f"Expected list length to be 2, but got {len(output)}"
+
+    if output[0].workflowUUID == wfUuid:
+        new_wfUuid = output[1].workflowUUID
+    else:
+        new_wfUuid = output[0].workflowUUID
+
+    info = _workflow_commands._get_workflow(config, new_wfUuid, True)
+    if info is not None:
+        assert info.status == "SUCCESS", f"Expected status to be SUCCESS"
+    else:
+        assert False, "Expected info to be not None"

{dbos-0.19.0a9 → dbos-0.20.0a2}/tests/test_fastapi.py

@@ -1,6 +1,8 @@
+import logging
 import uuid
 from typing import Tuple
 
+import pytest
 import sqlalchemy as sa
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
@@ -12,7 +14,9 @@ from dbos import DBOS
 from dbos._context import assert_current_dbos_context
 
 
-def test_simple_endpoint(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
+def test_simple_endpoint(
+    dbos_fastapi: Tuple[DBOS, FastAPI], caplog: pytest.LogCaptureFixture
+) -> None:
     dbos, app = dbos_fastapi
     client = TestClient(app)
 
@@ -32,6 +36,7 @@ def test_simple_endpoint(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
         res2 = test_step(var2)
         return res1 + res2
 
+    @app.get("/transaction/{var}")
     @DBOS.transaction()
     def test_transaction(var: str) -> str:
         rows = DBOS.sql_session.execute(sa.text("SELECT 1")).fetchall()
@@ -41,13 +46,27 @@ def test_simple_endpoint(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
     def test_step(var: str) -> str:
         return var
 
+    original_propagate = logging.getLogger("dbos").propagate
+    caplog.set_level(logging.WARNING, "dbos")
+    logging.getLogger("dbos").propagate = True
+
     response = client.get("/workflow/bob/bob")
     assert response.status_code == 200
     assert response.text == '"bob1bob"'
+    assert caplog.text == ""
 
     response = client.get("/endpoint/bob/bob")
     assert response.status_code == 200
     assert response.text == '"bob1bob"'
+    assert caplog.text == ""
+
+    response = client.get("/transaction/bob")
+    assert response.status_code == 200
+    assert response.text == '"bob1"'
+    assert caplog.text == ""
+
+    # Reset logging
+    logging.getLogger("dbos").propagate = original_propagate
 
 
 def test_start_workflow(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:

{dbos-0.19.0a9 → dbos-0.20.0a2}/tests/test_flask.py

@@ -1,6 +1,8 @@
+import logging
 import uuid
 from typing import Tuple
 
+import pytest
 import sqlalchemy as sa
 from flask import Flask, Response, jsonify
 
@@ -8,7 +10,9 @@ from dbos import DBOS
 from dbos._context import assert_current_dbos_context
 
 
-def test_flask_endpoint(dbos_flask: Tuple[DBOS, Flask]) -> None:
+def test_flask_endpoint(
+    dbos_flask: Tuple[DBOS, Flask], caplog: pytest.LogCaptureFixture
+) -> None:
     _, app = dbos_flask
 
     @app.route("/endpoint/<var1>/<var2>")
@@ -27,6 +31,7 @@ def test_flask_endpoint(dbos_flask: Tuple[DBOS, Flask]) -> None:
         result = res1 + res2
         return jsonify({"result": result})
 
+    @app.route("/transaction/<var>")
     @DBOS.transaction()
     def test_transaction(var: str) -> str:
         rows = DBOS.sql_session.execute(sa.text("SELECT 1")).fetchall()
@@ -39,13 +44,27 @@ def test_flask_endpoint(dbos_flask: Tuple[DBOS, Flask]) -> None:
     app.config["TESTING"] = True
     client = app.test_client()
 
+    original_propagate = logging.getLogger("dbos").propagate
+    caplog.set_level(logging.WARNING, "dbos")
+    logging.getLogger("dbos").propagate = True
+
     response = client.get("/endpoint/a/b")
     assert response.status_code == 200
     assert response.json == {"result": "a1b"}
+    assert caplog.text == ""
 
     response = client.get("/workflow/a/b")
     assert response.status_code == 200
     assert response.json == {"result": "a1b"}
+    assert caplog.text == ""
+
+    response = client.get("/transaction/bob")
+    assert response.status_code == 200
+    assert response.text == "bob1"
+    assert caplog.text == ""
+
+    # Reset logging
+    logging.getLogger("dbos").propagate = original_propagate
 
 
 def test_endpoint_recovery(dbos_flask: Tuple[DBOS, Flask]) -> None: