dbos 0.21.0a5__tar.gz → 0.22.0__tar.gz

{dbos-0.21.0a5 → dbos-0.22.0}/PKG-INFO

@@ -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 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.21.0a5 → dbos-0.22.0}/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-0.22.0/dbos/_recovery.py

@@ -0,0 +1,69 @@
+import os
+import threading
+import time
+import traceback
+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", 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(pending_workflows) > 0:
+        try:
+            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:
+            dbos.logger.error(
+                f"Exception encountered when recovering workflows: {traceback.format_exc()}"
+            )
+            raise e
+
+
+def recover_pending_workflows(
+    dbos: "DBOS", executor_ids: List[str] = ["local"]
+) -> List["WorkflowHandle[Any]"]:
+    workflow_handles: List["WorkflowHandle[Any]"] = []
+    for executor_id in executor_ids:
+        dbos.logger.debug(f"Recovering pending workflows for executor: {executor_id}")
+        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-0.21.0a5 → dbos-0.22.0}/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")