dbos 0.21.0a5__py3-none-any.whl → 0.22.0__py3-none-any.whl

dbos/_admin_server.py

@@ -16,6 +16,7 @@ if TYPE_CHECKING:
 _health_check_path = "/dbos-healthz"
 _workflow_recovery_path = "/dbos-workflow-recovery"
 _deactivate_path = "/deactivate"
+_workflow_queues_metadata_path = "/dbos-workflow-queues-metadata"
 # /workflows/:workflow_id/cancel
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
@@ -64,6 +65,26 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self.send_response(200)
             self._end_headers()
             self.wfile.write("deactivated".encode("utf-8"))
+        elif self.path == _workflow_queues_metadata_path:
+            queue_metadata_array = []
+            from ._dbos import _get_or_create_dbos_registry
+
+            registry = _get_or_create_dbos_registry()
+            for queue in registry.queue_info_map.values():
+                queue_metadata = {
+                    "name": queue.name,
+                    "concurrency": queue.concurrency,
+                    "workerConcurrency": queue.worker_concurrency,
+                    "rateLimit": queue.limiter,
+                }
+                # Remove keys with None values
+                queue_metadata = {
+                    k: v for k, v in queue_metadata.items() if v is not None
+                }
+                queue_metadata_array.append(queue_metadata)
+            self.send_response(200)
+            self._end_headers()
+            self.wfile.write(json.dumps(queue_metadata_array).encode("utf-8"))
         else:
             self.send_response(404)
             self._end_headers()

dbos/_cloudutils/cloudutils.py

@@ -29,9 +29,11 @@ class DBOSCloudCredentials:
 @dataclass
 class UserProfile:
     Name: str
-    Email: str
     Organization: str
-    SubscriptionPlan: str
+
+    def __init__(self, **kwargs: Any) -> None:
+        self.Name = kwargs.get("Name", "")
+        self.Organization = kwargs.get("Organization", "")
 
 
 class AppLanguages(Enum):

dbos/_cloudutils/databases.py

@@ -23,6 +23,10 @@ class UserDBCredentials:
     RoleName: str
     Password: str
 
+    def __init__(self, **kwargs: Any) -> None:
+        self.RoleName = kwargs.get("RoleName", "")
+        self.Password = kwargs.get("Password", "")
+
 
 @dataclass
 class UserDBInstance:

dbos/_context.py

@@ -49,7 +49,6 @@ class TracedAttributes(TypedDict, total=False):
 class DBOSContext:
     def __init__(self) -> None:
         self.executor_id = os.environ.get("DBOS__VMID", "local")
-        self.app_version = os.environ.get("DBOS__APPVERSION", "")
         self.app_id = os.environ.get("DBOS__APPID", "")
 
         self.logger = dbos_logger

dbos/_core.py

@@ -63,6 +63,7 @@ from ._registrations import (
     get_or_create_func_info,
     get_temp_workflow_type,
     set_dbos_func_name,
+    set_func_info,
     set_temp_workflow_type,
 )
 from ._roles import check_required_roles
@@ -162,7 +163,7 @@ def _init_workflow(
         "output": None,
         "error": None,
         "app_id": ctx.app_id,
-        "app_version": ctx.app_version,
+        "app_version": dbos.app_version,
         "executor_id": ctx.executor_id,
         "request": (
             _serialization.serialize(ctx.request) if ctx.request is not None else None
@@ -286,6 +287,7 @@ def execute_workflow_by_id(
         ctx.request = (
             _serialization.deserialize(request) if request is not None else None
         )
+        # If this function belongs to a configured class, add that class instance as its first argument
         if status["config_name"] is not None:
             config_name = status["config_name"]
             class_name = status["class_name"]
@@ -295,28 +297,9 @@ def execute_workflow_by_id(
                     workflow_id,
                     f"Cannot execute workflow because instance '{iname}' is not registered",
                 )
-
-            if startNew:
-                return start_workflow(
-                    dbos,
-                    wf_func,
-                    status["queue_name"],
-                    True,
-                    dbos._registry.instance_info_map[iname],
-                    *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"],
-                    )
+            class_instance = dbos._registry.instance_info_map[iname]
+            inputs["args"] = (class_instance,) + inputs["args"]
+        # If this function is a class method, add that class object as its first argument
         elif status["class_name"] is not None:
             class_name = status["class_name"]
             if class_name not in dbos._registry.class_info_map:
@@ -324,30 +307,20 @@ def execute_workflow_by_id(
                     workflow_id,
                     f"Cannot execute workflow because class '{class_name}' is not registered",
                 )
+            class_object = dbos._registry.class_info_map[class_name]
+            inputs["args"] = (class_object,) + inputs["args"]
 
-            if startNew:
-                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):
-                    return start_workflow(
-                        dbos,
-                        wf_func,
-                        status["queue_name"],
-                        True,
-                        dbos._registry.class_info_map[class_name],
-                        *inputs["args"],
-                        **inputs["kwargs"],
-                    )
+        if startNew:
+            return start_workflow(
+                dbos,
+                wf_func,
+                status["queue_name"],
+                True,
+                *inputs["args"],
+                **inputs["kwargs"],
+            )
         else:
-            if startNew:
+            with SetWorkflowID(workflow_id):
                 return start_workflow(
                     dbos,
                     wf_func,
@@ -356,16 +329,6 @@ def execute_workflow_by_id(
                     *inputs["args"],
                     **inputs["kwargs"],
                 )
-            else:
-                with SetWorkflowID(workflow_id):
-                    return start_workflow(
-                        dbos,
-                        wf_func,
-                        status["queue_name"],
-                        True,
-                        *inputs["args"],
-                        **inputs["kwargs"],
-                    )
 
 
 @overload
@@ -398,9 +361,12 @@ def start_workflow(
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> "WorkflowHandle[R]":
+    # If the function has a class, add the class object as its first argument
     fself: Optional[object] = None
     if hasattr(func, "__self__"):
         fself = func.__self__
+    if fself is not None:
+        args = (fself,) + args  # type: ignore
 
     fi = get_func_info(func)
     if fi is None:
@@ -436,17 +402,13 @@ def start_workflow(
     new_wf_ctx.id_assigned_for_next_workflow = new_wf_ctx.assign_workflow_id()
     new_wf_id = new_wf_ctx.id_assigned_for_next_workflow
 
-    gin_args: Tuple[Any, ...] = args
-    if fself is not None:
-        gin_args = (fself,)
-
     status = _init_workflow(
         dbos,
         new_wf_ctx,
         inputs=inputs,
         wf_name=get_dbos_func_name(func),
-        class_name=get_dbos_class_name(fi, func, gin_args),
-        config_name=get_config_name(fi, func, gin_args),
+        class_name=get_dbos_class_name(fi, func, args),
+        config_name=get_config_name(fi, func, args),
         temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
         max_recovery_attempts=fi.max_recovery_attempts,
@@ -464,27 +426,15 @@ def start_workflow(
         )
         return WorkflowHandlePolling(new_wf_id, dbos)
 
-    if fself is not None:
-        future = dbos._executor.submit(
-            cast(Callable[..., R], _execute_workflow_wthread),
-            dbos,
-            status,
-            func,
-            new_wf_ctx,
-            fself,
-            *args,
-            **kwargs,
-        )
-    else:
-        future = dbos._executor.submit(
-            cast(Callable[..., R], _execute_workflow_wthread),
-            dbos,
-            status,
-            func,
-            new_wf_ctx,
-            *args,
-            **kwargs,
-        )
+    future = dbos._executor.submit(
+        cast(Callable[..., R], _execute_workflow_wthread),
+        dbos,
+        status,
+        func,
+        new_wf_ctx,
+        *args,
+        **kwargs,
+    )
     return WorkflowHandleFuture(new_wf_id, future, dbos)
 
 
@@ -516,6 +466,8 @@ def workflow_wrapper(
 
     @wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> R:
+        fi = get_func_info(func)
+        assert fi is not None
         if dbosreg.dbos is None:
             raise DBOSException(
                 f"Function {func.__name__} invoked before DBOS initialized"
@@ -572,7 +524,7 @@ def decorate_workflow(
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     def _workflow_decorator(func: Callable[P, R]) -> Callable[P, R]:
         wrapped_func = workflow_wrapper(reg, func, max_recovery_attempts)
-        reg.register_wf_function(func.__qualname__, wrapped_func)
+        reg.register_wf_function(func.__qualname__, wrapped_func, "workflow")
         return wrapped_func
 
     return _workflow_decorator
@@ -724,8 +676,12 @@ def decorate_transaction(
         wrapped_wf = workflow_wrapper(dbosreg, temp_wf)
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "transaction")
-        dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
+        dbosreg.register_wf_function(
+            get_dbos_func_name(temp_wf), wrapped_wf, "transaction"
+        )
         wrapper.__orig_func = temp_wf  # type: ignore
+        set_func_info(wrapped_wf, get_or_create_func_info(func))
+        set_func_info(temp_wf, get_or_create_func_info(func))
 
         return cast(F, wrapper)
 
@@ -873,8 +829,10 @@ def decorate_step(
         wrapped_wf = workflow_wrapper(dbosreg, temp_wf)
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
         set_temp_workflow_type(temp_wf, "step")
-        dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
+        dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf, "step")
         wrapper.__orig_func = temp_wf  # type: ignore
+        set_func_info(wrapped_wf, get_or_create_func_info(func))
+        set_func_info(temp_wf, get_or_create_func_info(func))
 
         return cast(Callable[P, R], wrapper)
 

dbos/_db_wizard.py

@@ -1,7 +1,7 @@
 import json
 import os
 import time
-from typing import TYPE_CHECKING, Optional, TypedDict
+from typing import TYPE_CHECKING, Optional, TypedDict, cast
 
 import docker  # type: ignore
 import typer
@@ -45,8 +45,20 @@ def db_wizard(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
             f"Could not connect to Postgres: password authentication failed: {db_connection_error}"
         )
     db_config = config["database"]
+
+    # Read the config file and check if the database hostname/port/username are set. If so, skip the wizard.
+    with open(config_file_path, "r") as file:
+        content = file.read()
+        local_config = yaml.safe_load(content)
+        if "database" not in local_config:
+            local_config["database"] = {}
+        local_config = cast("ConfigFile", local_config)
+
     if (
-        db_config["hostname"] != "localhost"
+        local_config["database"].get("hostname")
+        or local_config["database"].get("port")
+        or local_config["database"].get("username")
+        or db_config["hostname"] != "localhost"
         or db_config["port"] != 5432
         or db_config["username"] != "postgres"
     ):

dbos/_dbos.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import asyncio
 import atexit
+import hashlib
+import inspect
 import json
 import os
 import sys
@@ -56,7 +58,7 @@ from ._registrations import (
 )
 from ._roles import default_required_roles, required_roles
 from ._scheduler import ScheduledWorkflow, scheduled
-from ._sys_db import WorkflowStatusString, reset_system_database
+from ._sys_db import reset_system_database
 from ._tracer import dbos_tracer
 
 if TYPE_CHECKING:
@@ -83,7 +85,11 @@ from ._context import (
     get_local_dbos_context,
 )
 from ._dbos_config import ConfigFile, load_config, set_env_vars
-from ._error import DBOSException, DBOSNonExistentWorkflowError
+from ._error import (
+    DBOSConflictingRegistrationError,
+    DBOSException,
+    DBOSNonExistentWorkflowError,
+)
 from ._logger import add_otlp_to_all_loggers, dbos_logger
 from ._sys_db import SystemDatabase
 
@@ -142,6 +148,7 @@ RegisteredJob = Tuple[
 class DBOSRegistry:
     def __init__(self) -> None:
         self.workflow_info_map: dict[str, Workflow[..., Any]] = {}
+        self.function_type_map: dict[str, str] = {}
         self.class_info_map: dict[str, type] = {}
         self.instance_info_map: dict[str, object] = {}
         self.queue_info_map: dict[str, Queue] = {}
@@ -149,7 +156,11 @@ class DBOSRegistry:
         self.dbos: Optional[DBOS] = None
         self.config: Optional[ConfigFile] = None
 
-    def register_wf_function(self, name: str, wrapped_func: F) -> None:
+    def register_wf_function(self, name: str, wrapped_func: F, functype: str) -> None:
+        if name in self.function_type_map:
+            if self.function_type_map[name] != functype:
+                raise DBOSConflictingRegistrationError(name)
+        self.function_type_map[name] = functype
         self.workflow_info_map[name] = wrapped_func
 
     def register_class(self, cls: type, ci: DBOSClassInfo) -> None:
@@ -186,6 +197,22 @@ class DBOSRegistry:
         else:
             self.instance_info_map[fn] = inst
 
+    def compute_app_version(self) -> str:
+        """
+        An application's version is computed from a hash of the source of its workflows.
+        This is guaranteed to be stable given identical source code because it uses an MD5 hash
+        and because it iterates through the workflows in sorted order.
+        This way, if the app's workflows are updated (which would break recovery), its version changes.
+        App version can be manually set through the DBOS__APPVERSION environment variable.
+        """
+        hasher = hashlib.md5()
+        sources = sorted(
+            [inspect.getsource(wf) for wf in self.workflow_info_map.values()]
+        )
+        for source in sources:
+            hasher.update(source.encode("utf-8"))
+        return hasher.hexdigest()
+
 
 class DBOS:
     """
@@ -283,6 +310,7 @@ class DBOS:
         self._executor_field: Optional[ThreadPoolExecutor] = None
         self._background_threads: List[threading.Thread] = []
         self._executor_id: str = os.environ.get("DBOS__VMID", "local")
+        self.app_version: str = os.environ.get("DBOS__APPVERSION", "")
 
         # If using FastAPI, set up middleware and lifecycle events
         if self.fastapi is not None:
@@ -305,7 +333,7 @@ class DBOS:
         temp_send_wf = workflow_wrapper(self._registry, send_temp_workflow)
         set_dbos_func_name(send_temp_workflow, TEMP_SEND_WF_NAME)
         set_temp_workflow_type(send_temp_workflow, "send")
-        self._registry.register_wf_function(TEMP_SEND_WF_NAME, temp_send_wf)
+        self._registry.register_wf_function(TEMP_SEND_WF_NAME, temp_send_wf, "send")
 
         for handler in dbos_logger.handlers:
             handler.flush()
@@ -351,6 +379,10 @@ class DBOS:
                 dbos_logger.warning(f"DBOS was already launched")
                 return
             self._launched = True
+            if self.app_version == "":
+                self.app_version = self._registry.compute_app_version()
+            dbos_logger.info(f"Application version: {self.app_version}")
+            dbos_tracer.app_version = self.app_version
             self._executor_field = ThreadPoolExecutor(max_workers=64)
             self._sys_db_field = SystemDatabase(self.config)
             self._app_db_field = ApplicationDatabase(self.config)
@@ -359,9 +391,19 @@ class DBOS:
                 admin_port = 3001
             self._admin_server_field = AdminServer(dbos=self, port=admin_port)
 
-            if not os.environ.get("DBOS__VMID"):
-                workflow_ids = self._sys_db.get_pending_workflows("local")
-                self._executor.submit(startup_recovery_thread, self, workflow_ids)
+            workflow_ids = self._sys_db.get_pending_workflows(
+                self._executor_id, self.app_version
+            )
+            if (len(workflow_ids)) > 0:
+                self.logger.info(
+                    f"Recovering {len(workflow_ids)} workflows from application version {self.app_version}"
+                )
+            else:
+                self.logger.info(
+                    f"No workflows to recover from application version {self.app_version}"
+                )
+
+            self._executor.submit(startup_recovery_thread, self, workflow_ids)
 
             # Listen to notifications
             notification_listener_thread = threading.Thread(
@@ -398,13 +440,13 @@ class DBOS:
                 self._background_threads.append(poller_thread)
             self._registry.pollers = []
 
-            dbos_logger.info("DBOS launched")
+            dbos_logger.info("DBOS launched!")
 
             # Flush handlers and add OTLP to all loggers if enabled
             # to enable their export in DBOS Cloud
             for handler in dbos_logger.handlers:
                 handler.flush()
-            add_otlp_to_all_loggers()
+            add_otlp_to_all_loggers(self.app_version)
         except Exception:
             dbos_logger.error(f"DBOS failed to launch: {traceback.format_exc()}")
             raise
@@ -613,6 +655,7 @@ class DBOS:
             workflow_id=workflow_id,
             status=stat["status"],
             name=stat["name"],
+            executor_id=stat["executor_id"],
             recovery_attempts=stat["recovery_attempts"],
             class_name=stat["class_name"],
             config_name=stat["config_name"],
@@ -909,6 +952,7 @@ class WorkflowStatus:
         workflow_id(str):  The ID of the workflow execution
         status(str):  The status of the execution, from `WorkflowStatusString`
         name(str): The workflow function name
+        executor_id(str): The ID of the executor running the workflow
         class_name(str): For member functions, the name of the class containing the workflow function
         config_name(str): For instance member functions, the name of the class instance for the execution
         queue_name(str): For workflows that are or were queued, the queue name
@@ -922,6 +966,7 @@ class WorkflowStatus:
     workflow_id: str
     status: str
     name: str
+    executor_id: Optional[str]
     class_name: Optional[str]
     config_name: Optional[str]
     queue_name: Optional[str]
@@ -994,6 +1039,10 @@ def _dbos_exit_hook() -> None:
         )
         return
     if not _dbos_global_instance._launched:
+        if _dbos_global_instance.fastapi is not None:
+            # FastAPI lifespan middleware will call launch/destroy, so we can ignore this.
+            # This is likely to happen during fastapi dev runs, where the reloader loads the module multiple times.
+            return
         print("DBOS exiting; DBOS exists but launch() was not called")
         dbos_logger.warning("DBOS exiting; DBOS exists but launch() was not called")
         return

dbos/_error.py

@@ -36,6 +36,7 @@ class DBOSErrorCode(Enum):
     MaxStepRetriesExceeded = 7
     NotAuthorized = 8
     ConflictingWorkflowError = 9
+    ConflictingRegistrationError = 25
 
 
 class DBOSWorkflowConflictIDError(DBOSException):
@@ -127,3 +128,13 @@ class DBOSMaxStepRetriesExceeded(DBOSException):
             "Step reached maximum retries.",
             dbos_error_code=DBOSErrorCode.MaxStepRetriesExceeded.value,
         )
+
+
+class DBOSConflictingRegistrationError(DBOSException):
+    """Exception raised when conflicting decorators are applied to the same function."""
+
+    def __init__(self, name: str) -> None:
+        super().__init__(
+            f"Operation (Name: {name}) is already registered with a conflicting function type",
+            dbos_error_code=DBOSErrorCode.ConflictingRegistrationError.value,
+        )

dbos/_logger.py

@@ -86,8 +86,9 @@ def config_logger(config: "ConfigFile") -> None:
         dbos_logger.addFilter(_otlp_transformer)
 
 
-def add_otlp_to_all_loggers() -> None:
+def add_otlp_to_all_loggers(app_version: str) -> None:
     if _otlp_handler is not None and _otlp_transformer is not None:
+        _otlp_transformer.app_version = app_version
         root = logging.root
 
         root.addHandler(_otlp_handler)

dbos/_queue.py

@@ -76,7 +76,9 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
                     execute_workflow_by_id(dbos, id)
             except OperationalError as e:
                 # Ignore serialization error
-                if not isinstance(e.orig, errors.SerializationFailure):
+                if not isinstance(
+                    e.orig, (errors.SerializationFailure, errors.LockNotAvailable)
+                ):
                     dbos.logger.warning(
                         f"Exception encountered in queue thread: {traceback.format_exc()}"
                     )

dbos/_recovery.py

@@ -6,20 +6,29 @@ from typing import TYPE_CHECKING, Any, List
 
 from ._core import execute_workflow_by_id
 from ._error import DBOSWorkflowFunctionNotFoundError
+from ._sys_db import GetPendingWorkflowsOutput
 
 if TYPE_CHECKING:
     from ._dbos import DBOS, WorkflowHandle
 
 
-def startup_recovery_thread(dbos: "DBOS", workflow_ids: List[str]) -> None:
+def startup_recovery_thread(
+    dbos: "DBOS", pending_workflows: List[GetPendingWorkflowsOutput]
+) -> None:
     """Attempt to recover local pending workflows on startup using a background thread."""
     stop_event = threading.Event()
     dbos.stop_events.append(stop_event)
-    while not stop_event.is_set() and len(workflow_ids) > 0:
+    while not stop_event.is_set() and len(pending_workflows) > 0:
         try:
-            for workflowID in list(workflow_ids):
-                execute_workflow_by_id(dbos, workflowID)
-                workflow_ids.remove(workflowID)
+            for pending_workflow in list(pending_workflows):
+                if (
+                    pending_workflow.queue_name
+                    and pending_workflow.queue_name != "_dbos_internal_queue"
+                ):
+                    dbos._sys_db.clear_queue_assignment(pending_workflow.workflow_uuid)
+                    continue
+                execute_workflow_by_id(dbos, pending_workflow.workflow_uuid)
+                pending_workflows.remove(pending_workflow)
         except DBOSWorkflowFunctionNotFoundError:
             time.sleep(1)
         except Exception as e:
@@ -34,17 +43,27 @@ def recover_pending_workflows(
 ) -> List["WorkflowHandle[Any]"]:
     workflow_handles: List["WorkflowHandle[Any]"] = []
     for executor_id in executor_ids:
-        if executor_id == "local" and os.environ.get("DBOS__VMID"):
-            dbos.logger.debug(
-                f"Skip local recovery because it's running in a VM: {os.environ.get('DBOS__VMID')}"
-            )
         dbos.logger.debug(f"Recovering pending workflows for executor: {executor_id}")
-        workflow_ids = dbos._sys_db.get_pending_workflows(executor_id)
-        dbos.logger.debug(f"Pending workflows: {workflow_ids}")
-
-        for workflowID in workflow_ids:
-            handle = execute_workflow_by_id(dbos, workflowID)
-            workflow_handles.append(handle)
-
-    dbos.logger.info("Recovered pending workflows")
+        pending_workflows = dbos._sys_db.get_pending_workflows(
+            executor_id, dbos.app_version
+        )
+        for pending_workflow in pending_workflows:
+            if (
+                pending_workflow.queue_name
+                and pending_workflow.queue_name != "_dbos_internal_queue"
+            ):
+                try:
+                    dbos._sys_db.clear_queue_assignment(pending_workflow.workflow_uuid)
+                    workflow_handles.append(
+                        dbos.retrieve_workflow(pending_workflow.workflow_uuid)
+                    )
+                except Exception as e:
+                    dbos.logger.error(e)
+            else:
+                workflow_handles.append(
+                    execute_workflow_by_id(dbos, pending_workflow.workflow_uuid)
+                )
+        dbos.logger.info(
+            f"Recovering {len(pending_workflows)} workflows from version {dbos.app_version}"
+        )
     return workflow_handles

dbos/_registrations.py

@@ -1,4 +1,5 @@
 import inspect
+from dataclasses import dataclass
 from enum import Enum
 from types import FunctionType
 from typing import Any, Callable, List, Literal, Optional, Tuple, Type, cast
@@ -31,9 +32,9 @@ def set_temp_workflow_type(f: Any, name: TempWorkflowType) -> None:
     setattr(f, "dbos_temp_workflow_type", name)
 
 
+@dataclass
 class DBOSClassInfo:
-    def __init__(self) -> None:
-        self.def_required_roles: Optional[List[str]] = None
+    def_required_roles: Optional[List[str]] = None
 
 
 class DBOSFuncType(Enum):
@@ -44,12 +45,12 @@ class DBOSFuncType(Enum):
     Instance = 4
 
 
+@dataclass
 class DBOSFuncInfo:
-    def __init__(self) -> None:
-        self.class_info: Optional[DBOSClassInfo] = None
-        self.func_type: DBOSFuncType = DBOSFuncType.Unknown
-        self.required_roles: Optional[List[str]] = None
-        self.max_recovery_attempts = DEFAULT_MAX_RECOVERY_ATTEMPTS
+    class_info: Optional[DBOSClassInfo] = None
+    func_type: DBOSFuncType = DBOSFuncType.Unknown
+    required_roles: Optional[List[str]] = None
+    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS
 
 
 def get_or_create_class_info(cls: Type[Any]) -> DBOSClassInfo:
@@ -110,6 +111,10 @@ def get_or_create_func_info(func: Callable[..., Any]) -> DBOSFuncInfo:
     return fi
 
 
+def set_func_info(func: Callable[..., Any], fi: DBOSFuncInfo) -> None:
+    setattr(func, "dbos_func_decorator_info", fi)
+
+
 def get_class_info(cls: Type[Any]) -> Optional[DBOSClassInfo]:
     if hasattr(cls, "dbos_class_decorator_info"):
         ci: DBOSClassInfo = getattr(cls, "dbos_class_decorator_info")

dbos/_sys_db.py

@@ -1,4 +1,5 @@
 import datetime
+import logging
 import os
 import re
 import threading
@@ -13,6 +14,7 @@ from typing import (
     Optional,
     Sequence,
     Set,
+    Tuple,
     TypedDict,
     cast,
 )
@@ -22,8 +24,8 @@ 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 sqlalchemy.sql import func
 
 from . import _serialization
 from ._dbos_config import ConfigFile
@@ -140,6 +142,12 @@ class GetWorkflowsOutput:
         self.workflow_uuids = workflow_uuids
 
 
+class GetPendingWorkflowsOutput:
+    def __init__(self, *, workflow_uuid: str, queue_name: Optional[str] = None):
+        self.workflow_uuid: str = workflow_uuid
+        self.queue_name: Optional[str] = queue_name
+
+
 class WorkflowInformation(TypedDict, total=False):
     workflow_uuid: str
     status: WorkflowStatuses  # The status of the workflow.
@@ -183,6 +191,10 @@ class SystemDatabase:
             host=config["database"]["hostname"],
             port=config["database"]["port"],
             database="postgres",
+            # fills the "application_name" column in pg_stat_activity
+            query={
+                "application_name": f"dbos_transact_{os.environ.get('DBOS__VMID', 'local')}"
+            },
         )
         engine = sa.create_engine(postgres_db_url)
         with engine.connect() as conn:
@@ -201,6 +213,10 @@ class SystemDatabase:
             host=config["database"]["hostname"],
             port=config["database"]["port"],
             database=sysdb_name,
+            # fills the "application_name" column in pg_stat_activity
+            query={
+                "application_name": f"dbos_transact_{os.environ.get('DBOS__VMID', 'local')}"
+            },
         )
 
         # Create a connection pool for the system database
@@ -214,6 +230,7 @@ class SystemDatabase:
         )
         alembic_cfg = Config()
         alembic_cfg.set_main_option("script_location", migration_dir)
+        logging.getLogger("alembic").setLevel(logging.WARNING)
         # Alembic requires the % in URL-escaped parameters to itself be escaped to %%.
         escaped_conn_string = re.sub(
             r"%(?=[0-9A-Fa-f]{2})",
@@ -221,7 +238,12 @@ class SystemDatabase:
             self.engine.url.render_as_string(hide_password=False),
         )
         alembic_cfg.set_main_option("sqlalchemy.url", escaped_conn_string)
-        command.upgrade(alembic_cfg, "head")
+        try:
+            command.upgrade(alembic_cfg, "head")
+        except Exception as e:
+            dbos_logger.warning(
+                f"Exception during system database construction. This is most likely because the system database was configured using a later version of DBOS: {e}"
+            )
 
         self.notification_conn: Optional[psycopg.connection.Connection] = None
         self.notifications_map: Dict[str, threading.Condition] = {}
@@ -288,6 +310,7 @@ class SystemDatabase:
                     recovery_attempts=(
                         SystemSchema.workflow_status.c.recovery_attempts + 1
                     ),
+                    updated_at=func.extract("epoch", func.now()) * 1000,
                 ),
             )
         )
@@ -385,6 +408,7 @@ class SystemDatabase:
                     status=status["status"],
                     output=status["output"],
                     error=status["error"],
+                    updated_at=func.extract("epoch", func.now()) * 1000,
                 ),
             )
         )
@@ -465,6 +489,7 @@ class SystemDatabase:
                     SystemSchema.workflow_status.c.authenticated_roles,
                     SystemSchema.workflow_status.c.assumed_role,
                     SystemSchema.workflow_status.c.queue_name,
+                    SystemSchema.workflow_status.c.executor_id,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -479,7 +504,7 @@ class SystemDatabase:
                 "error": None,
                 "app_id": None,
                 "app_version": None,
-                "executor_id": None,
+                "executor_id": row[10],
                 "request": row[2],
                 "recovery_attempts": row[3],
                 "authenticated_user": row[6],
@@ -665,7 +690,7 @@ class SystemDatabase:
 
     def get_workflows(self, input: GetWorkflowsInput) -> GetWorkflowsOutput:
         query = sa.select(SystemSchema.workflow_status.c.workflow_uuid).order_by(
-            SystemSchema.workflow_status.c.created_at.desc()
+            SystemSchema.workflow_status.c.created_at.asc()
         )
         if input.name:
             query = query.where(SystemSchema.workflow_status.c.name == input.name)
@@ -711,7 +736,7 @@ class SystemDatabase:
                 SystemSchema.workflow_queue.c.workflow_uuid
                 == SystemSchema.workflow_status.c.workflow_uuid,
             )
-            .order_by(SystemSchema.workflow_status.c.created_at.desc())
+            .order_by(SystemSchema.workflow_status.c.created_at.asc())
         )
 
         if input.get("name"):
@@ -746,16 +771,29 @@ class SystemDatabase:
 
         return GetWorkflowsOutput(workflow_uuids)
 
-    def get_pending_workflows(self, executor_id: str) -> list[str]:
+    def get_pending_workflows(
+        self, executor_id: str, app_version: str
+    ) -> list[GetPendingWorkflowsOutput]:
         with self.engine.begin() as c:
             rows = c.execute(
-                sa.select(SystemSchema.workflow_status.c.workflow_uuid).where(
+                sa.select(
+                    SystemSchema.workflow_status.c.workflow_uuid,
+                    SystemSchema.workflow_status.c.queue_name,
+                ).where(
                     SystemSchema.workflow_status.c.status
                     == WorkflowStatusString.PENDING.value,
                     SystemSchema.workflow_status.c.executor_id == executor_id,
+                    SystemSchema.workflow_status.c.application_version == app_version,
                 )
             ).fetchall()
-            return [row[0] for row in rows]
+
+            return [
+                GetPendingWorkflowsOutput(
+                    workflow_uuid=row.workflow_uuid,
+                    queue_name=row.queue_name,
+                )
+                for row in rows
+            ]
 
     def record_operation_result(
         self, result: OperationResultInternal, conn: Optional[sa.Connection] = None
@@ -1284,6 +1322,55 @@ class SystemDatabase:
             # 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.
+
+            # First lets figure out how many tasks the worker can dequeue
+            running_tasks_query = (
+                sa.select(
+                    SystemSchema.workflow_queue.c.executor_id,
+                    sa.func.count().label("task_count"),
+                )
+                .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
+                .where(
+                    SystemSchema.workflow_queue.c.executor_id.isnot(
+                        None
+                    )  # Task is dequeued
+                )
+                .where(
+                    SystemSchema.workflow_queue.c.completed_at_epoch_ms.is_(
+                        None
+                    )  # Task is not completed
+                )
+                .group_by(SystemSchema.workflow_queue.c.executor_id)
+            )
+            running_tasks_result = c.execute(running_tasks_query).fetchall()
+            running_tasks_result_dict = {row[0]: row[1] for row in running_tasks_result}
+            running_tasks_for_this_worker = running_tasks_result_dict.get(
+                executor_id, 0
+            )  # Get count for current executor
+
+            max_tasks = float("inf")
+            if queue.worker_concurrency is not None:
+                # Worker local concurrency limit should always be >= running_tasks_for_this_worker
+                # This should never happen but a check + warning doesn't hurt
+                if running_tasks_for_this_worker > queue.worker_concurrency:
+                    dbos_logger.warning(
+                        f"Number of tasks on this worker ({running_tasks_for_this_worker}) exceeds the worker concurrency limit ({queue.worker_concurrency})"
+                    )
+                max_tasks = max(
+                    0, queue.worker_concurrency - running_tasks_for_this_worker
+                )
+            if queue.concurrency is not None:
+                total_running_tasks = sum(running_tasks_result_dict.values())
+                # Queue global concurrency limit should always be >= running_tasks_count
+                # This should never happen but a check + warning doesn't hurt
+                if total_running_tasks > queue.concurrency:
+                    dbos_logger.warning(
+                        f"Total running tasks ({total_running_tasks}) exceeds the global concurrency limit ({queue.concurrency})"
+                    )
+                available_tasks = max(0, queue.concurrency - total_running_tasks)
+                max_tasks = min(max_tasks, available_tasks)
+
+            # Lookup tasks
             query = (
                 sa.select(
                     SystemSchema.workflow_queue.c.workflow_uuid,
@@ -1292,29 +1379,25 @@ class SystemDatabase:
                 )
                 .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,
-                    )
-                )
+                .where(SystemSchema.workflow_queue.c.executor_id == None)
                 .order_by(SystemSchema.workflow_queue.c.created_at_epoch_ms.asc())
+                .with_for_update(nowait=True)  # Error out early
             )
-            # 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)
+            # Apply limit only if max_tasks is finite
+            if max_tasks != float("inf"):
+                query = query.limit(int(max_tasks))
 
             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]
+            # Get the workflow IDs
+            dequeued_ids: List[str] = [row[0] for row in rows]
+            if len(dequeued_ids) > 0:
+                dbos_logger.debug(
+                    f"[{queue.name}] dequeueing {len(dequeued_ids)} task(s)"
+                )
             ret_ids: list[str] = []
-            dbos_logger.debug(f"[{queue.name}] dequeueing {len(dequeued_ids)} task(s)")
-            for id in dequeued_ids:
 
+            for id in dequeued_ids:
                 # If we have a limiter, stop starting functions when the number
                 # of functions started this period exceeds the limit.
                 if queue.limiter is not None:
@@ -1375,6 +1458,19 @@ class SystemDatabase:
                     .values(completed_at_epoch_ms=int(time.time() * 1000))
                 )
 
+    def clear_queue_assignment(self, workflow_id: str) -> None:
+        with self.engine.begin() as c:
+            c.execute(
+                sa.update(SystemSchema.workflow_queue)
+                .where(SystemSchema.workflow_queue.c.workflow_uuid == workflow_id)
+                .values(executor_id=None, started_at_epoch_ms=None)
+            )
+            c.execute(
+                sa.update(SystemSchema.workflow_status)
+                .where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
+                .values(executor_id=None, status=WorkflowStatusString.ENQUEUED.value)
+            )
+
 
 def reset_system_database(config: ConfigFile) -> None:
     sysdb_name = (

{dbos-0.21.0a5.dist-info → dbos-0.22.0.dist-info}/METADATA

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

{dbos-0.21.0a5.dist-info → dbos-0.22.0.dist-info}/RECORD

@@ -1,26 +1,26 @@
-dbos-0.21.0a5.dist-info/METADATA,sha256=1LAPN_Eu4X7ptR9xM02sR1tryFaQ42KeKGkeOm3yCQQ,5309
-dbos-0.21.0a5.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-dbos-0.21.0a5.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-0.21.0a5.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.22.0.dist-info/METADATA,sha256=K4tMreHMHoB2VORi7cfuh1oUiOezJIy7EHVsr5Zh0fA,5307
+dbos-0.22.0.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+dbos-0.22.0.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-0.22.0.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=CxRHBHEthPL4PZoLbZhp3rdm44-KkRTT2-7DkK9d4QQ,724
-dbos/_admin_server.py,sha256=PJgneZG9-64TapZrPeJtt73puAswRImCE5uce2k2PKU,4750
+dbos/_admin_server.py,sha256=YiVn5lywz2Vg8_juyNHOYl0HVEy48--7b4phwK7r92o,5732
 dbos/_app_db.py,sha256=_tv2vmPjjiaikwgxH3mqxgJ4nUUcG2-0uMXKWCqVu1c,5509
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
 dbos/_cloudutils/authentication.py,sha256=V0fCWQN9stCkhbuuxgPTGpvuQcDqfU3KAxPAh01vKW4,5007
-dbos/_cloudutils/cloudutils.py,sha256=5e3CW1deSW-dI5G3QN0XbiVsBhyqT8wu7fuV2f8wtGU,7688
-dbos/_cloudutils/databases.py,sha256=x4187Djsyoa-QaG3Kog8JT2_GERsnqa93LIVanmVUmg,8393
-dbos/_context.py,sha256=FHB_fpE4fQt4fIJvAmMMsbY4xHwH77gsW01cFsRZjsE,17779
-dbos/_core.py,sha256=nGiXyYgV8H5TRRZG0e8HCd5IZimufYQLmKNr7nBbwbo,36564
+dbos/_cloudutils/cloudutils.py,sha256=YC7jGsIopT0KveLsqbRpQk2KlRBk-nIRC_UCgep4f3o,7797
+dbos/_cloudutils/databases.py,sha256=_shqaqSvhY4n2ScgQ8IP5PDZvzvcx3YBKV8fj-cxhSY,8543
+dbos/_context.py,sha256=gikN5lUVqnvR-ISoOElXYeYsR_BO2whebB3YP2DJBM4,17713
+dbos/_core.py,sha256=b1IL3LDRGmomHhcs929n_G2pfFSmNuSxFrwhcBuo20k,35519
 dbos/_croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
-dbos/_db_wizard.py,sha256=xgKLna0_6Xi50F3o8msRosXba8NScHlpJR5ICVCkHDQ,7534
-dbos/_dbos.py,sha256=y5RgXPxdNsnketphphGMlaYZABFEQpr78UK-Xyja6dk,36216
+dbos/_db_wizard.py,sha256=6tfJaCRa1NtkUdNW75a2yvi_mEgnPJ9C1HP2zPG1hCU,8067
+dbos/_dbos.py,sha256=h25S5Mjl1JAfwMEpqRdyBSuUB_HI3TC8J9Nqtqy_XwQ,38453
 dbos/_dbos_config.py,sha256=DfiqVVxNqnafkocSzLqBp1Ig5vCviDTDK_GO3zTtQqI,8298
-dbos/_error.py,sha256=vtaSsG0QW6cRlwfZ4zzZWy_IHCZlomwSlrDyGWuyn8c,4337
+dbos/_error.py,sha256=NqlobQneZ2ycCQacXc8a38TIOHxFRjBXdF40i3wZUaA,4775
 dbos/_fastapi.py,sha256=ke03vqsSYDnO6XeOtOVFXj0-f-v1MGsOxa9McaROvNc,3616
 dbos/_flask.py,sha256=DZKUZR5-xOzPI7tYZ53r2PvvHVoAb8SYwLzMVFsVfjI,2608
 dbos/_kafka.py,sha256=o6DbwnsYRDtvVTZVsN7BAK8cdP79AfoWX3Q7CGY2Yuo,4199
 dbos/_kafka_message.py,sha256=NYvOXNG3Qn7bghn1pv3fg4Pbs86ILZGcK4IB-MLUNu0,409
-dbos/_logger.py,sha256=iYwbA7DLyXalWa2Yu07HO6Xm301nRuenMU64GgwUMkU,3576
+dbos/_logger.py,sha256=hNEeOgR9yOwdgcOuvnW_wN9rbfpTk5OowPNhEJmjoQE,3644
 dbos/_migrations/env.py,sha256=38SIGVbmn_VV2x2u1aHLcPOoWgZ84eCymf3g_NljmbU,1626
 dbos/_migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
 dbos/_migrations/versions/04ca4f231047_workflow_queues_executor_id.py,sha256=ICLPl8CN9tQXMsLDsAj8z1TsL831-Z3F8jSBvrR-wyw,736
@@ -31,9 +31,9 @@ dbos/_migrations/versions/d76646551a6b_job_queue_limiter.py,sha256=8PyFi8rd6CN-m
 dbos/_migrations/versions/d76646551a6c_workflow_queue.py,sha256=G942nophZ2uC2vc4hGBC02Ptng1715roTjY3xiyzZU4,729
 dbos/_migrations/versions/eab0cc1d9a14_job_queue.py,sha256=uvhFOtqbBreCePhAxZfIT0qCAI7BiZTou9wt6QnbY7c,1412
 dbos/_outcome.py,sha256=FDMgWVjZ06vm9xO-38H17mTqBImUYQxgKs_bDCSIAhE,6648
-dbos/_queue.py,sha256=o_aczwualJTMoXb0XXL-Y5QH77OEukWzuerogbWi2ho,2779
-dbos/_recovery.py,sha256=ehruOebVRl8qHftydStGZd-VjJ7p9oYtYXDwyT4ZHRY,1874
-dbos/_registrations.py,sha256=mei6q6_3R5uei8i_Wo_TqGZs85s10shOekDX41sFYD0,6642
+dbos/_queue.py,sha256=eZiapBcyn70-viW0y9fo7u09V6_VF5ACNGJxD-U_dNM,2844
+dbos/_recovery.py,sha256=GtNMvPFM9qetob-gCU9FPI2fo_BGZYRws4EFSpOuFa4,2675
+dbos/_registrations.py,sha256=_zy6k944Ll8QwqU12Kr3OP23ukVtm8axPNN1TS_kJRc,6717
 dbos/_request.py,sha256=cX1B3Atlh160phgS35gF1VEEV4pD126c9F3BDgBmxZU,929
 dbos/_roles.py,sha256=iOsgmIAf1XVzxs3gYWdGRe1B880YfOw5fpU7Jwx8_A8,2271
 dbos/_scheduler.py,sha256=0I3e8Y-OIBG3wiUCIskShd-Sk_eUFCFyRB5u4L7IHXI,1940
@@ -41,7 +41,7 @@ dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=KeyoPrF7hy_ODXV7QNike_VFSD74QBRfQ76D7QyE9HI,966
 dbos/_schemas/system_database.py,sha256=rwp4EvCSaXcUoMaRczZCvETCxGp72k3-hvLyGUDkih0,5163
 dbos/_serialization.py,sha256=YCYv0qKAwAZ1djZisBC7khvKqG-5OcIv9t9EC5PFIog,1743
-dbos/_sys_db.py,sha256=DLoSddcraHPUMiA5zWIck9Kes7_UO39RL2CRPqtYyz0,58247
+dbos/_sys_db.py,sha256=aGa3KDKUVeK9d7r_yK21uoKV3LJGzRzsg8S3Bea5j_U,62685
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/dbos-db-starter/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/dbos-db-starter/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
@@ -60,4 +60,4 @@ dbos/cli/cli.py,sha256=_tXw2IQrWW7fV_h51f_R99vEBSi6aMLz-vCOxKaENiQ,14155
 dbos/dbos-config.schema.json,sha256=X5TpXNcARGceX0zQs0fVgtZW_Xj9uBbY5afPt9Rz9yk,5741
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.21.0a5.dist-info/RECORD,,
+dbos-0.22.0.dist-info/RECORD,,