Improve Errors For Unregistered Functions and Classes (#429)

Author: kraftp

dbos/_context.py

@@ -753,3 +753,34 @@ def __exit__(
             assert ctx.assumed_role == self.assume_role
         ctx.assumed_role = self.prior_role
         return False  # Did not handle
+
+
+class UseLogAttributes:
+    """Temporarily set context attributes for logging"""
+
+    def __init__(self, *, workflow_id: str = "") -> None:
+        self.workflow_id = workflow_id
+        self.created_ctx = False
+
+    def __enter__(self) -> UseLogAttributes:
+        ctx = get_local_dbos_context()
+        if ctx is None:
+            self.created_ctx = True
+            _set_local_dbos_context(DBOSContext())
+        ctx = assert_current_dbos_context()
+        self.saved_workflow_id = ctx.workflow_id
+        ctx.workflow_id = self.workflow_id
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        ctx = assert_current_dbos_context()
+        ctx.workflow_id = self.saved_workflow_id
+        # Clean up the basic context if we created it
+        if self.created_ctx:
+            _clear_local_dbos_context()
+        return False  # Did not handle

dbos/_core.py

@@ -5,7 +5,6 @@
 import sys
 import threading
 import time
-import traceback
 from concurrent.futures import Future
 from functools import wraps
 from typing import (
@@ -66,7 +65,6 @@
     get_dbos_func_name,
     get_func_info,
     get_or_create_func_info,
-    get_temp_workflow_type,
     set_dbos_func_name,
     set_func_info,
     set_temp_workflow_type,
@@ -452,7 +450,7 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
     if not wf_func:
         raise DBOSWorkflowFunctionNotFoundError(
             workflow_id,
-            f"Cannot execute workflow because {status['name']} is not a registered workflow function",
+            f"{status['name']} is not a registered workflow function",
         )
     with DBOSContextEnsure():
         # If this function belongs to a configured class, add that class instance as its first argument
@@ -463,7 +461,7 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
             if iname not in dbos._registry.instance_info_map:
                 raise DBOSWorkflowFunctionNotFoundError(
                     workflow_id,
-                    f"Cannot execute workflow because instance '{iname}' is not registered",
+                    f"configured class instance '{iname}' is not registered",
                 )
             class_instance = dbos._registry.instance_info_map[iname]
             inputs["args"] = (class_instance,) + inputs["args"]
@@ -473,7 +471,7 @@ def execute_workflow_by_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[An
             if class_name not in dbos._registry.class_info_map:
                 raise DBOSWorkflowFunctionNotFoundError(
                     workflow_id,
-                    f"Cannot execute workflow because class '{class_name}' is not registered",
+                    f"class '{class_name}' is not registered",
                 )
             class_object = dbos._registry.class_info_map[class_name]
             inputs["args"] = (class_object,) + inputs["args"]
@@ -534,7 +532,7 @@ def start_workflow(
     if fi is None:
         raise DBOSWorkflowFunctionNotFoundError(
             "<NONE>",
-            f"start_workflow: function {func.__name__} is not registered",
+            f"{func.__name__} is not a registered workflow function",
         )
 
     func = cast("Workflow[P, R]", func.__orig_func)  # type: ignore
@@ -630,7 +628,7 @@ async def start_workflow_async(
     if fi is None:
         raise DBOSWorkflowFunctionNotFoundError(
             "<NONE>",
-            f"start_workflow: function {func.__name__} is not registered",
+            f"{func.__name__} is not a registered workflow function",
         )
 
     func = cast("Workflow[P, R]", func.__orig_func)  # type: ignore

dbos/_dbos.py

@@ -211,6 +211,10 @@ def register_poller(
     def register_instance(self, inst: object) -> None:
         config_name = getattr(inst, "config_name")
         class_name = _class_fqn(inst.__class__)
+        if self.dbos and self.dbos._launched:
+            dbos_logger.warning(
+                f"Configured instance {config_name} of class {class_name} was registered after DBOS was launched. This may cause errors during workflow recovery. All configured instances should be instantiated before DBOS is launched."
+            )
         fn = f"{class_name}/{config_name}"
         if fn in self.instance_info_map:
             if self.instance_info_map[fn] is not inst:

dbos/_error.py

@@ -106,7 +106,7 @@ class DBOSWorkflowFunctionNotFoundError(DBOSException):
 
     def __init__(self, workflow_id: str, message: Optional[str] = None):
         super().__init__(
-            f"Workflow function not found for workflow ID {workflow_id}: {message}",
+            f"Could not execute workflow {workflow_id}: {message}",
             dbos_error_code=DBOSErrorCode.WorkflowFunctionNotFound.value,
         )
 

dbos/_recovery.py

@@ -2,6 +2,7 @@
 import time
 from typing import TYPE_CHECKING, Any, List
 
+from dbos._context import UseLogAttributes
 from dbos._utils import GlobalParams
 
 from ._core import execute_workflow_by_id
@@ -29,17 +30,19 @@ def startup_recovery_thread(
     stop_event = threading.Event()
     dbos.background_thread_stop_events.append(stop_event)
     while not stop_event.is_set() and len(pending_workflows) > 0:
-        try:
-            for pending_workflow in list(pending_workflows):
+        for pending_workflow in list(pending_workflows):
+            try:
                 _recover_workflow(dbos, pending_workflow)
                 pending_workflows.remove(pending_workflow)
-        except DBOSWorkflowFunctionNotFoundError:
-            time.sleep(1)
-        except Exception as e:
-            dbos.logger.error(
-                f"Exception encountered when recovering workflows:", exc_info=e
-            )
-            raise
+            except DBOSWorkflowFunctionNotFoundError:
+                time.sleep(1)
+            except Exception as e:
+                with UseLogAttributes(workflow_id=pending_workflow.workflow_uuid):
+                    dbos.logger.error(
+                        f"Exception encountered when recovering workflow {pending_workflow.workflow_uuid}:",
+                        exc_info=e,
+                    )
+                raise
 
 
 def recover_pending_workflows(
@@ -56,9 +59,11 @@ def recover_pending_workflows(
                 handle = _recover_workflow(dbos, pending_workflow)
                 workflow_handles.append(handle)
             except Exception as e:
-                dbos.logger.error(
-                    f"Exception encountered when recovering workflows:", exc_info=e
-                )
+                with UseLogAttributes(workflow_id=pending_workflow.workflow_uuid):
+                    dbos.logger.error(
+                        f"Exception encountered when recovering workflow {pending_workflow.workflow_uuid}:",
+                        exc_info=e,
+                    )
                 raise
         dbos.logger.info(
             f"Recovering {len(pending_workflows)} workflows for executor {executor_id} from version {GlobalParams.app_version}"

dbos/_registrations.py

@@ -13,7 +13,7 @@ def get_dbos_func_name(f: Any) -> str:
     if hasattr(f, "dbos_function_name"):
         return str(getattr(f, "dbos_function_name"))
     raise DBOSWorkflowFunctionNotFoundError(
-        "<NONE>", f"function {f.__name__} is not registered"
+        "<NONE>", f"{f.__name__} is not a registered workflow function"
     )
 
 

tests/test_failures.py

@@ -8,12 +8,14 @@
 from sqlalchemy.exc import InvalidRequestError, OperationalError
 
 from dbos import DBOS, Queue, SetWorkflowID
+from dbos._dbos_config import DBOSConfig
 from dbos._error import (
     DBOSAwaitedWorkflowCancelledError,
     DBOSMaxStepRetriesExceeded,
     DBOSNotAuthorizedError,
     DBOSQueueDeduplicatedError,
     DBOSUnexpectedStepError,
+    DBOSWorkflowFunctionNotFoundError,
     MaxRecoveryAttemptsExceededError,
 )
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
@@ -499,3 +501,24 @@ def __init__(self, one: int, two: int) -> None:
     assert output is None
     assert isinstance(exception, str)
     assert "Message: 1, 2" in exception
+
+
+def test_unregistered_workflow(dbos: DBOS, config: DBOSConfig) -> None:
+
+    @DBOS.workflow()
+    def workflow() -> None:
+        return
+
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
+        workflow()
+
+    dbos._sys_db.update_workflow_outcome(wfid, "PENDING")
+
+    DBOS.destroy(destroy_registry=True)
+    config["executor_id"] = str(uuid.uuid4())
+    DBOS(config=config)
+    DBOS.launch()
+
+    with pytest.raises(DBOSWorkflowFunctionNotFoundError):
+        DBOS._recover_pending_workflows()