dbos 0.19.0a4__py3-none-any.whl → 0.20.0__py3-none-any.whl

dbos/_sys_db.py

@@ -28,6 +28,7 @@ from sqlalchemy.exc import DBAPIError
 from . import _serialization
 from ._dbos_config import ConfigFile
 from ._error import (
+    DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
     DBOSException,
     DBOSNonExistentWorkflowError,
@@ -288,8 +289,14 @@ class SystemDatabase:
                 ),
             )
         else:
-            cmd = cmd.on_conflict_do_nothing()
-        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status)  # 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)
@@ -297,37 +304,53 @@ class SystemDatabase:
             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]
-                wf_status = row[1]
-                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:
@@ -344,7 +367,7 @@ class SystemDatabase:
         with self.engine.begin() as c:
             stmt = (
                 sa.update(SystemSchema.workflow_status)
-                .where(SystemSchema.workflow_inputs.c.workflow_uuid == workflow_uuid)
+                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
                 .values(
                     status=status,
                 )
@@ -356,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)
                 )
@@ -538,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
@@ -582,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)
@@ -1233,3 +1265,46 @@ class SystemDatabase:
                     .where(SystemSchema.workflow_queue.c.workflow_uuid == workflow_id)
                     .values(completed_at_epoch_ms=int(time.time() * 1000))
                 )
+
+
+def reset_system_database(config: ConfigFile) -> None:
+    sysdb_name = (
+        config["database"]["sys_db_name"]
+        if "sys_db_name" in config["database"] and config["database"]["sys_db_name"]
+        else config["database"]["app_db_name"] + SystemSchema.sysdb_suffix
+    )
+    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",
+    )
+    try:
+        # Connect to postgres default database
+        engine = sa.create_engine(postgres_db_url)
+
+        with engine.connect() as conn:
+            # Set autocommit required for database dropping
+            conn.execution_options(isolation_level="AUTOCOMMIT")
+
+            # Terminate existing connections
+            conn.execute(
+                sa.text(
+                    """
+                SELECT pg_terminate_backend(pg_stat_activity.pid)
+                FROM pg_stat_activity
+                WHERE pg_stat_activity.datname = :db_name
+                AND pid <> pg_backend_pid()
+            """
+                ),
+                {"db_name": sysdb_name},
+            )
+
+            # Drop the database
+            conn.execute(sa.text(f"DROP DATABASE IF EXISTS {sysdb_name}"))
+
+    except sa.exc.SQLAlchemyError as e:
+        dbos_logger.error(f"Error resetting system database: {str(e)}")
+        raise e

dbos/_workflow_commands.py

@@ -0,0 +1,171 @@
+import importlib
+import os
+import sys
+from typing import Any, List, Optional, cast
+
+import typer
+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,
+    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 _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/_github_init.py

@@ -0,0 +1,107 @@
+import os
+from base64 import b64decode
+from typing import List, TypedDict
+
+import requests
+
+DEMO_REPO_API = "https://api.github.com/repos/dbos-inc/dbos-demo-apps"
+PY_DEMO_PATH = "python/"
+BRANCH = "main"
+
+
+class GitHubTreeItem(TypedDict):
+    path: str
+    mode: str
+    type: str
+    sha: str
+    url: str
+    size: int
+
+
+class GitHubTree(TypedDict):
+    sha: str
+    url: str
+    tree: List[GitHubTreeItem]
+    truncated: bool
+
+
+class GitHubItem(TypedDict):
+    sha: str
+    node_id: str
+    url: str
+    content: str
+    encoding: str
+    size: int
+
+
+def _fetch_github(url: str) -> requests.Response:
+    headers = {}
+    github_token = os.getenv("GITHUB_TOKEN")
+    if github_token:
+        headers["Authorization"] = f"Bearer {github_token}"
+
+    response = requests.get(url, headers=headers)
+
+    if not response.ok:
+        if response.headers.get("x-ratelimit-remaining") == "0":
+            raise Exception(
+                "Error fetching from GitHub API: rate limit exceeded.\n"
+                "Please wait a few minutes and try again.\n"
+                "To increase the limit, you can create a personal access token and set it in the GITHUB_TOKEN environment variable.\n"
+                "Details: https://docs.github.com/en/rest/using-the-rest-api/rate-limits-for-the-rest-api"
+            )
+        elif response.status_code == 401:
+            raise Exception(
+                f"Error fetching content from GitHub {url}: {response.status_code} {response.reason}.\n"
+                "Please ensure your GITHUB_TOKEN environment variable is set to a valid personal access token."
+            )
+        raise Exception(
+            f"Error fetching content from GitHub {url}: {response.status_code} {response.reason}"
+        )
+
+    return response
+
+
+def _fetch_github_tree(tag: str) -> List[GitHubTreeItem]:
+    response = _fetch_github(f"{DEMO_REPO_API}/git/trees/{tag}?recursive=1")
+    tree_data: GitHubTree = response.json()
+    return tree_data["tree"]
+
+
+def _fetch_github_item(url: str) -> str:
+    response = _fetch_github(url)
+    item: GitHubItem = response.json()
+    return b64decode(item["content"]).decode("utf-8")
+
+
+def create_template_from_github(app_name: str, template_name: str) -> None:
+    print(
+        f"Creating a new application named {app_name} from the template {template_name}"
+    )
+
+    tree = _fetch_github_tree(BRANCH)
+    template_path = f"{PY_DEMO_PATH}{template_name}/"
+
+    files_to_download = [
+        item
+        for item in tree
+        if item["path"].startswith(template_path) and item["type"] == "blob"
+    ]
+
+    # Download every file from the template
+    for item in files_to_download:
+        raw_content = _fetch_github_item(item["url"])
+        file_path = item["path"].replace(template_path, "")
+        target_path = os.path.join(".", file_path)
+
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+        # Write file with proper permissions
+        with open(target_path, "w", encoding="utf-8") as f:
+            f.write(raw_content)
+        os.chmod(target_path, int(item["mode"], 8))
+
+    print(
+        f"Downloaded {len(files_to_download)} files from the template GitHub repository"
+    )

dbos/cli/_template_init.py

@@ -0,0 +1,98 @@
+import os
+import shutil
+import typing
+from os import path
+from typing import Any
+
+import tomlkit
+from rich import print
+
+
+def get_templates_directory() -> str:
+    import dbos
+
+    package_dir = path.abspath(path.dirname(dbos.__file__))
+    return path.join(package_dir, "_templates")
+
+
+def _copy_dbos_template(src: str, dst: str, ctx: dict[str, str]) -> None:
+    with open(src, "r") as f:
+        content = f.read()
+
+    for key, value in ctx.items():
+        content = content.replace(f"${{{key}}}", value)
+
+    with open(dst, "w") as f:
+        f.write(content)
+
+
+def _copy_template_dir(src_dir: str, dst_dir: str, ctx: dict[str, str]) -> None:
+
+    for root, dirs, files in os.walk(src_dir, topdown=True):
+        dirs[:] = [d for d in dirs if d != "__package"]
+
+        dst_root = path.join(dst_dir, path.relpath(root, src_dir))
+        if len(dirs) == 0:
+            os.makedirs(dst_root, exist_ok=True)
+        else:
+            for dir in dirs:
+                os.makedirs(path.join(dst_root, dir), exist_ok=True)
+
+        for file in files:
+            src = path.join(root, file)
+            base, ext = path.splitext(file)
+
+            dst = path.join(dst_root, base if ext == ".dbos" else file)
+            if path.exists(dst):
+                print(f"[yellow]File {dst} already exists, skipping[/yellow]")
+                continue
+
+            if ext == ".dbos":
+                _copy_dbos_template(src, dst, ctx)
+            else:
+                shutil.copy(src, dst)
+
+
+def copy_template(src_dir: str, project_name: str, config_mode: bool) -> None:
+
+    dst_dir = path.abspath(".")
+
+    package_name = project_name.replace("-", "_")
+    ctx = {
+        "project_name": project_name,
+        "package_name": package_name,
+        "migration_command": "alembic upgrade head",
+    }
+
+    if config_mode:
+        ctx["package_name"] = "."
+        ctx["migration_command"] = "echo 'No migrations specified'"
+        _copy_dbos_template(
+            os.path.join(src_dir, "dbos-config.yaml.dbos"),
+            os.path.join(dst_dir, "dbos-config.yaml"),
+            ctx,
+        )
+    else:
+        _copy_template_dir(src_dir, dst_dir, ctx)
+        _copy_template_dir(
+            path.join(src_dir, "__package"), path.join(dst_dir, package_name), ctx
+        )
+
+
+def get_project_name() -> typing.Union[str, None]:
+    name = None
+    try:
+        with open("pyproject.toml", "rb") as file:
+            pyproj = typing.cast(dict[str, Any], tomlkit.load(file))
+            name = typing.cast(str, pyproj["project"]["name"])
+    except:
+        pass
+
+    if name == None:
+        try:
+            _, parent = path.split(path.abspath("."))
+            name = parent
+        except:
+            pass
+
+    return name