dbos 0.18.0a1__py3-none-any.whl → 0.19.0__py3-none-any.whl

dbos/_sys_db.py

@@ -13,7 +13,6 @@ from typing import (
     Optional,
     Sequence,
     Set,
-    Tuple,
     TypedDict,
     cast,
 )
@@ -23,12 +22,15 @@ import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 from alembic import command
 from alembic.config import Config
+from sqlalchemy import or_
 from sqlalchemy.exc import DBAPIError
 
 from . import _serialization
 from ._dbos_config import ConfigFile
 from ._error import (
+    DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
+    DBOSException,
     DBOSNonExistentWorkflowError,
     DBOSWorkflowConflictIDError,
 )
@@ -249,7 +251,9 @@ class SystemDatabase:
         *,
         conn: Optional[sa.Connection] = None,
         max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
-    ) -> None:
+    ) -> WorkflowStatuses:
+        wf_status: WorkflowStatuses = status["status"]
+
         cmd = pg.insert(SystemSchema.workflow_status).values(
             workflow_uuid=status["workflow_uuid"],
             status=status["status"],
@@ -285,49 +289,75 @@ class SystemDatabase:
                 ),
             )
         else:
-            cmd = cmd.on_conflict_do_nothing()
-        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts)  # type: ignore
+            # A blank update so that we can return the existing status
+            cmd = cmd.on_conflict_do_update(
+                index_elements=["workflow_uuid"],
+                set_=dict(
+                    recovery_attempts=SystemSchema.workflow_status.c.recovery_attempts
+                ),
+            )
+        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
 
         if conn is not None:
             results = conn.execute(cmd)
         else:
             with self.engine.begin() as c:
                 results = c.execute(cmd)
-        if in_recovery:
-            row = results.fetchone()
-            if row is not None:
-                recovery_attempts: int = row[0]
-                if recovery_attempts > max_recovery_attempts:
-                    with self.engine.begin() as c:
-                        c.execute(
-                            sa.delete(SystemSchema.workflow_queue).where(
-                                SystemSchema.workflow_queue.c.workflow_uuid
-                                == status["workflow_uuid"]
-                            )
+
+        row = results.fetchone()
+        if row is not None:
+            # Check the started workflow matches the expected name, class_name, config_name, and queue_name
+            # A mismatch indicates a workflow starting with the same UUID but different functions, which would throw an exception.
+            recovery_attempts: int = row[0]
+            wf_status = row[1]
+            err_msg: Optional[str] = None
+            if row[2] != status["name"]:
+                err_msg = f"Workflow already exists with a different function name: {row[2]}, but the provided function name is: {status['name']}"
+            elif row[3] != status["class_name"]:
+                err_msg = f"Workflow already exists with a different class name: {row[3]}, but the provided class name is: {status['class_name']}"
+            elif row[4] != status["config_name"]:
+                err_msg = f"Workflow already exists with a different config name: {row[4]}, but the provided config name is: {status['config_name']}"
+            elif row[5] != status["queue_name"]:
+                # This is a warning because a different queue name is not necessarily an error.
+                dbos_logger.warning(
+                    f"Workflow already exists in queue: {row[5]}, but the provided queue name is: {status['queue_name']}. The queue is not updated."
+                )
+            if err_msg is not None:
+                raise DBOSConflictingWorkflowError(status["workflow_uuid"], err_msg)
+
+            if in_recovery and recovery_attempts > max_recovery_attempts:
+                with self.engine.begin() as c:
+                    c.execute(
+                        sa.delete(SystemSchema.workflow_queue).where(
+                            SystemSchema.workflow_queue.c.workflow_uuid
+                            == status["workflow_uuid"]
+                        )
+                    )
+                    c.execute(
+                        sa.update(SystemSchema.workflow_status)
+                        .where(
+                            SystemSchema.workflow_status.c.workflow_uuid
+                            == status["workflow_uuid"]
+                        )
+                        .where(
+                            SystemSchema.workflow_status.c.status
+                            == WorkflowStatusString.PENDING.value
                         )
-                        c.execute(
-                            sa.update(SystemSchema.workflow_status)
-                            .where(
-                                SystemSchema.workflow_status.c.workflow_uuid
-                                == status["workflow_uuid"]
-                            )
-                            .where(
-                                SystemSchema.workflow_status.c.status
-                                == WorkflowStatusString.PENDING.value
-                            )
-                            .values(
-                                status=WorkflowStatusString.RETRIES_EXCEEDED.value,
-                                queue_name=None,
-                            )
+                        .values(
+                            status=WorkflowStatusString.RETRIES_EXCEEDED.value,
+                            queue_name=None,
                         )
-                    raise DBOSDeadLetterQueueError(
-                        status["workflow_uuid"], max_recovery_attempts
                     )
+                raise DBOSDeadLetterQueueError(
+                    status["workflow_uuid"], max_recovery_attempts
+                )
 
         # Record we have exported status for this single-transaction workflow
         if status["workflow_uuid"] in self._temp_txn_wf_ids:
             self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
 
+        return wf_status
+
     def set_workflow_status(
         self,
         workflow_uuid: str,
@@ -349,7 +379,7 @@ class SystemDatabase:
                 stmt = (
                     sa.update(SystemSchema.workflow_status)
                     .where(
-                        SystemSchema.workflow_inputs.c.workflow_uuid == workflow_uuid
+                        SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid
                     )
                     .values(recovery_attempts=reset_recovery_attempts)
                 )
@@ -405,7 +435,10 @@ class SystemDatabase:
                     res["output"]
                 )
                 return resstat
-            return None
+            else:
+                raise DBOSException(
+                    "Workflow status record not found. This should not happen! \033[1m Hint: Check if your workflow is deterministic.\033[0m"
+                )
         stat = self.get_workflow_status(workflow_uuid)
         self.record_operation_result(
             {
@@ -528,18 +561,27 @@ class SystemDatabase:
                 workflow_uuid=workflow_uuid,
                 inputs=inputs,
             )
-            .on_conflict_do_nothing()
+            .on_conflict_do_update(
+                index_elements=["workflow_uuid"],
+                set_=dict(workflow_uuid=SystemSchema.workflow_inputs.c.workflow_uuid),
+            )
+            .returning(SystemSchema.workflow_inputs.c.inputs)
         )
         if conn is not None:
-            conn.execute(cmd)
+            row = conn.execute(cmd).fetchone()
         else:
             with self.engine.begin() as c:
-                c.execute(cmd)
-
+                row = c.execute(cmd).fetchone()
+        if row is not None and row[0] != inputs:
+            dbos_logger.warning(
+                f"Workflow inputs for {workflow_uuid} changed since the first call! Use the original inputs."
+            )
+            # TODO: actually changing the input
         if workflow_uuid in self._temp_txn_wf_ids:
             # Clean up the single-transaction tracking sets
             self._exported_temp_txn_wf_status.discard(workflow_uuid)
             self._temp_txn_wf_ids.discard(workflow_uuid)
+        return
 
     def get_workflow_inputs(
         self, workflow_uuid: str
@@ -572,12 +614,12 @@ class SystemDatabase:
         if input.start_time:
             query = query.where(
                 SystemSchema.workflow_status.c.created_at
-                >= datetime.datetime.fromisoformat(input.start_time).timestamp()
+                >= datetime.datetime.fromisoformat(input.start_time).timestamp() * 1000
             )
         if input.end_time:
             query = query.where(
                 SystemSchema.workflow_status.c.created_at
-                <= datetime.datetime.fromisoformat(input.end_time).timestamp()
+                <= datetime.datetime.fromisoformat(input.end_time).timestamp() * 1000
             )
         if input.status:
             query = query.where(SystemSchema.workflow_status.c.status == input.status)
@@ -1130,27 +1172,38 @@ class SystemDatabase:
                 if num_recent_queries >= queue.limiter["limit"]:
                     return []
 
-            # Select not-yet-completed functions in the queue ordered by the
-            # time at which they were enqueued.
-            # If there is a concurrency limit N, select only the N most recent
+            # Dequeue functions eligible for this worker and ordered by the time at which they were enqueued.
+            # If there is a global or local concurrency limit N, select only the N oldest enqueued
             # functions, else select all of them.
             query = (
                 sa.select(
                     SystemSchema.workflow_queue.c.workflow_uuid,
                     SystemSchema.workflow_queue.c.started_at_epoch_ms,
+                    SystemSchema.workflow_queue.c.executor_id,
                 )
                 .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
                 .where(SystemSchema.workflow_queue.c.completed_at_epoch_ms == None)
+                .where(
+                    # Only select functions that have not been started yet or have been started by this worker
+                    or_(
+                        SystemSchema.workflow_queue.c.executor_id == None,
+                        SystemSchema.workflow_queue.c.executor_id == executor_id,
+                    )
+                )
                 .order_by(SystemSchema.workflow_queue.c.created_at_epoch_ms.asc())
             )
-            if queue.concurrency is not None:
+            # Set a dequeue limit if necessary
+            if queue.worker_concurrency is not None:
+                query = query.limit(queue.worker_concurrency)
+            elif queue.concurrency is not None:
                 query = query.limit(queue.concurrency)
 
-            # From the functions retrieved, get the workflow IDs of the functions
-            # that have not yet been started so we can start them.
             rows = c.execute(query).fetchall()
+
+            # Now, get the workflow IDs of functions that have not yet been started
             dequeued_ids: List[str] = [row[0] for row in rows if row[1] is None]
             ret_ids: list[str] = []
+            dbos_logger.debug(f"[{queue.name}] dequeueing {len(dequeued_ids)} task(s)")
             for id in dequeued_ids:
 
                 # If we have a limiter, stop starting functions when the number
@@ -1173,11 +1226,11 @@ class SystemDatabase:
                     )
                 )
 
-                # Then give it a start time
+                # Then give it a start time and assign the executor ID
                 c.execute(
                     SystemSchema.workflow_queue.update()
                     .where(SystemSchema.workflow_queue.c.workflow_uuid == id)
-                    .values(started_at_epoch_ms=start_time_ms)
+                    .values(started_at_epoch_ms=start_time_ms, executor_id=executor_id)
                 )
                 ret_ids.append(id)
 

dbos/_templates/hello/dbos-config.yaml.dbos

@@ -9,10 +9,6 @@ runtimeConfig:
   start:
     - "fastapi run ${package_name}/main.py"
 database:
-  hostname: localhost
-  port: 5432
-  username: postgres
-  password: ${PGPASSWORD}
   migrate:
     - ${migration_command}
 telemetry:

dbos/_workflow_commands.py

@@ -0,0 +1,172 @@
+from typing import Any, List, Optional, cast
+
+import typer
+from rich import print
+
+from dbos import DBOS
+
+from . import _serialization, load_config
+from ._dbos_config import ConfigFile, _is_valid_app_name
+from ._sys_db import (
+    GetWorkflowsInput,
+    GetWorkflowsOutput,
+    SystemDatabase,
+    WorkflowStatuses,
+    WorkflowStatusInternal,
+    WorkflowStatusString,
+)
+
+
+class WorkflowInformation:
+    workflowUUID: str
+    status: WorkflowStatuses
+    workflowName: str
+    workflowClassName: Optional[str]
+    workflowConfigName: Optional[str]
+    input: Optional[_serialization.WorkflowInputs]  # JSON (jsonpickle)
+    output: Optional[str]  # JSON (jsonpickle)
+    error: Optional[str]  # JSON (jsonpickle)
+    executor_id: Optional[str]
+    app_version: Optional[str]
+    app_id: Optional[str]
+    request: Optional[str]  # JSON (jsonpickle)
+    recovery_attempts: Optional[int]
+    authenticated_user: Optional[str]
+    assumed_role: Optional[str]
+    authenticated_roles: Optional[str]  # JSON list of roles.
+    queue_name: Optional[str]
+
+
+def _list_workflows(
+    config: ConfigFile,
+    li: int,
+    user: Optional[str],
+    starttime: Optional[str],
+    endtime: Optional[str],
+    status: Optional[str],
+    request: bool,
+    appversion: Optional[str],
+) -> List[WorkflowInformation]:
+
+    sys_db = None
+
+    try:
+        sys_db = SystemDatabase(config)
+
+        input = GetWorkflowsInput()
+        input.authenticated_user = user
+        input.start_time = starttime
+        input.end_time = endtime
+        if status is not None:
+            input.status = cast(WorkflowStatuses, status)
+        input.application_version = appversion
+        input.limit = li
+
+        output: GetWorkflowsOutput = sys_db.get_workflows(input)
+
+        infos: List[WorkflowInformation] = []
+
+        if output.workflow_uuids is None:
+            typer.echo("No workflows found")
+            return {}
+
+        for workflow_id in output.workflow_uuids:
+            info = _get_workflow_info(
+                sys_db, workflow_id, request
+            )  # Call the method for each ID
+
+            if info is not None:
+                infos.append(info)
+
+        return infos
+    except Exception as e:
+        typer.echo(f"Error listing workflows: {e}")
+        return []
+    finally:
+        if sys_db:
+            sys_db.destroy()
+
+
+def _get_workflow(
+    config: ConfigFile, uuid: str, request: bool
+) -> Optional[WorkflowInformation]:
+    sys_db = None
+
+    try:
+        sys_db = SystemDatabase(config)
+
+        info = _get_workflow_info(sys_db, uuid, request)
+        return info
+
+    except Exception as e:
+        typer.echo(f"Error getting workflow: {e}")
+        return None
+    finally:
+        if sys_db:
+            sys_db.destroy()
+
+
+def _cancel_workflow(config: ConfigFile, uuid: str) -> None:
+    # config = load_config()
+    sys_db = None
+
+    try:
+        sys_db = SystemDatabase(config)
+        sys_db.set_workflow_status(uuid, WorkflowStatusString.CANCELLED, False)
+        return
+
+    except Exception as e:
+        typer.echo(f"Failed to connect to DBOS system database: {e}")
+        return None
+    finally:
+        if sys_db:
+            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]:
+
+    info = sys_db.get_workflow_status(workflowUUID)
+    if info is None:
+        return None
+
+    winfo = WorkflowInformation()
+
+    winfo.workflowUUID = workflowUUID
+    winfo.status = info["status"]
+    winfo.workflowName = info["name"]
+    winfo.workflowClassName = info["class_name"]
+    winfo.workflowConfigName = info["config_name"]
+    winfo.executor_id = info["executor_id"]
+    winfo.app_version = info["app_version"]
+    winfo.app_id = info["app_id"]
+    winfo.recovery_attempts = info["recovery_attempts"]
+    winfo.authenticated_user = info["authenticated_user"]
+    winfo.assumed_role = info["assumed_role"]
+    winfo.authenticated_roles = info["authenticated_roles"]
+    winfo.queue_name = info["queue_name"]
+
+    # no input field
+    input_data = sys_db.get_workflow_inputs(workflowUUID)
+    if input_data is not None:
+        winfo.input = input_data
+
+    if info.get("status") == "SUCCESS":
+        result = sys_db.await_workflow_result(workflowUUID)
+        winfo.output = result
+    elif info.get("status") == "ERROR":
+        try:
+            sys_db.await_workflow_result(workflowUUID)
+        except Exception as e:
+            winfo.error = str(e)
+
+    if not getRequest:
+        winfo.request = None
+
+    return winfo

dbos/cli.py

@@ -8,6 +8,7 @@ import typing
 from os import path
 from typing import Any
 
+import jsonpickle  # type: ignore
 import sqlalchemy as sa
 import tomlkit
 import typer
@@ -17,12 +18,21 @@ from typing_extensions import Annotated
 
 from dbos._schemas.system_database import SystemSchema
 
-from . import load_config
+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,
+)
 
 app = typer.Typer()
+workflow = typer.Typer()
+
+app.add_typer(workflow, name="workflow", help="Manage DBOS workflows")
 
 
 def _on_windows() -> bool:
@@ -333,5 +343,94 @@ def reset(
             sys_db.destroy()
 
 
+@workflow.command(help="List workflows for your application")
+def list(
+    limit: Annotated[
+        int,
+        typer.Option("--limit", "-l", help="Limit the results returned"),
+    ] = 10,
+    user: Annotated[
+        typing.Optional[str],
+        typer.Option("--user", "-u", help="Retrieve workflows run by this user"),
+    ] = None,
+    starttime: Annotated[
+        typing.Optional[str],
+        typer.Option(
+            "--start-time",
+            "-s",
+            help="Retrieve workflows starting after this timestamp (ISO 8601 format)",
+        ),
+    ] = None,
+    endtime: Annotated[
+        typing.Optional[str],
+        typer.Option(
+            "--end-time",
+            "-e",
+            help="Retrieve workflows starting before this timestamp (ISO 8601 format)",
+        ),
+    ] = None,
+    status: Annotated[
+        typing.Optional[str],
+        typer.Option(
+            "--status",
+            "-S",
+            help="Retrieve workflows with this status (PENDING, SUCCESS, ERROR, RETRIES_EXCEEDED, ENQUEUED, or CANCELLED)",
+        ),
+    ] = None,
+    appversion: Annotated[
+        typing.Optional[str],
+        typer.Option(
+            "--application-version",
+            "-v",
+            help="Retrieve workflows with this application version",
+        ),
+    ] = None,
+    request: Annotated[
+        bool,
+        typer.Option("--request", help="Retrieve workflow request information"),
+    ] = True,
+    appdir: Annotated[
+        typing.Optional[str],
+        typer.Option("--app-dir", "-d", help="Specify the application root directory"),
+    ] = None,
+) -> None:
+    config = load_config()
+    workflows = _list_workflows(
+        config, limit, user, starttime, endtime, status, request, appversion
+    )
+    print(jsonpickle.encode(workflows, unpicklable=False))
+
+
+@workflow.command(help="Retrieve the status of a workflow")
+def get(
+    uuid: Annotated[str, typer.Argument()],
+    appdir: Annotated[
+        typing.Optional[str],
+        typer.Option("--app-dir", "-d", help="Specify the application root directory"),
+    ] = None,
+    request: Annotated[
+        bool,
+        typer.Option("--request", help="Retrieve workflow request information"),
+    ] = True,
+) -> None:
+    config = load_config()
+    print(jsonpickle.encode(_get_workflow(config, uuid, request), unpicklable=False))
+
+
+@workflow.command(
+    help="Cancel a workflow so it is no longer automatically retried or restarted"
+)
+def cancel(
+    uuid: Annotated[str, typer.Argument()],
+    appdir: Annotated[
+        typing.Optional[str],
+        typer.Option("--app-dir", "-d", help="Specify the application root directory"),
+    ] = None,
+) -> None:
+    config = load_config()
+    _cancel_workflow(config, uuid)
+    print(f"Workflow {uuid} has been cancelled")
+
+
 if __name__ == "__main__":
     app()

dbos/dbos-config.schema.json

@@ -81,13 +81,7 @@
             "type": "array",
             "description": "Specify a list of user DB rollback commands to run"
           }
-        },
-        "required": [
-          "hostname",
-          "port",
-          "username",
-          "password"
-        ]
+        }
       },
       "telemetry": {
         "type": "object",
@@ -181,9 +175,6 @@
         "type": "string",
         "deprecated": true
       }
-    },
-    "required": [
-      "database"
-    ]
+    }
   }