Check for conflicting decorators (#220)

Author: chuck-dbos

dbos/_core.py

@@ -524,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
@@ -676,7 +676,9 @@ def temp_wf(*args: Any, **kwargs: Any) -> Any:
         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))
@@ -827,7 +829,7 @@ async def temp_wf_async(*args: Any, **kwargs: Any) -> Any:
         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))

dbos/_dbos.py

@@ -85,7 +85,11 @@
     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
 
@@ -144,14 +148,19 @@ def _get_or_create_dbos_registry() -> DBOSRegistry:
 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] = {}
         self.pollers: list[RegisteredJob] = []
         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:
@@ -324,7 +333,7 @@ def send_temp_workflow(
         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()

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 @@ def __init__(self) -> None:
             "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,
+        )

tests/test_dbos.py

@@ -16,7 +16,7 @@
 
 # Private API because this is a test
 from dbos._context import assert_current_dbos_context, get_local_dbos_context
-from dbos._error import DBOSMaxStepRetriesExceeded
+from dbos._error import DBOSConflictingRegistrationError, DBOSMaxStepRetriesExceeded
 from dbos._schemas.system_database import SystemSchema
 from dbos._sys_db import GetWorkflowsInput
 
@@ -1208,6 +1208,20 @@ def test_workflow(var: str) -> str:
     assert test_workflow(var) == var
 
 
+def test_double_decoration(dbos: DBOS) -> None:
+    with pytest.raises(
+        DBOSConflictingRegistrationError,
+        match="is already registered with a conflicting function type",
+    ):
+
+        @DBOS.step()
+        @DBOS.transaction()
+        def my_function() -> None:
+            pass
+
+        my_function()
+
+
 def test_app_version(config: ConfigFile) -> None:
     def is_hex(s: str) -> bool:
         return all(c in "0123456789abcdefABCDEF" for c in s)