dbos 0.8.0a0__py3-none-any.whl → 0.8.0a7__py3-none-any.whl

dbos/__init__.py

@@ -1,5 +1,5 @@
 from . import error as error
-from .context import DBOSContextEnsure, SetWorkflowID
+from .context import DBOSContextEnsure, DBOSContextSetAuth, SetWorkflowID
 from .dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
 from .dbos_config import ConfigFile, get_dbos_database_url, load_config
 from .kafka_message import KafkaMessage
@@ -11,6 +11,7 @@ __all__ = [
     "DBOS",
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",
+    "DBOSContextSetAuth",
     "GetWorkflowsInput",
     "KafkaMessage",
     "SetWorkflowID",

dbos/context.py

@@ -492,6 +492,39 @@ class EnterDBOSHandler:
         return False  # Did not handle
 
 
+class DBOSContextSetAuth(DBOSContextEnsure):
+    def __init__(self, user: Optional[str], roles: Optional[List[str]]) -> None:
+        self.created_ctx = False
+        self.user = user
+        self.roles = roles
+        self.prev_user: Optional[str] = None
+        self.prev_roles: Optional[List[str]] = None
+
+    def __enter__(self) -> DBOSContext:
+        ctx = get_local_dbos_context()
+        if ctx is None:
+            self.created_ctx = True
+            set_local_dbos_context(DBOSContext())
+        ctx = assert_current_dbos_context()
+        self.prev_user = ctx.authenticated_user
+        self.prev_roles = ctx.authenticated_roles
+        ctx.set_authentication(self.user, self.roles)
+        return ctx
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        ctx = assert_current_dbos_context()
+        ctx.set_authentication(self.prev_user, self.prev_roles)
+        # Clean up the basic context if we created it
+        if self.created_ctx:
+            clear_local_dbos_context()
+        return False  # Did not handle
+
+
 class DBOSAssumeRole:
     def __init__(self, assume_role: Optional[str]) -> None:
         self.prior_role: Optional[str] = None

dbos/core.py

@@ -4,17 +4,7 @@ import time
 import traceback
 from concurrent.futures import Future
 from functools import wraps
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Generic,
-    List,
-    Optional,
-    Tuple,
-    TypeVar,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Tuple, TypeVar, cast
 
 from dbos.application_database import ApplicationDatabase, TransactionResultInternal
 
@@ -48,6 +38,7 @@ from dbos.error import (
     DBOSWorkflowFunctionNotFoundError,
 )
 from dbos.registrations import (
+    DEFAULT_MAX_RECOVERY_ATTEMPTS,
     get_config_name,
     get_dbos_class_name,
     get_dbos_func_name,
@@ -61,10 +52,10 @@ from dbos.roles import check_required_roles
 from dbos.system_database import (
     GetEventWorkflowContext,
     OperationResultInternal,
-    WorkflowInputs,
     WorkflowStatusInternal,
     WorkflowStatusString,
 )
+from dbos.utils import WorkflowInputs
 
 if TYPE_CHECKING:
     from dbos.dbos import DBOS, Workflow, WorkflowHandle, WorkflowStatus, _DBOSRegistry
@@ -128,6 +119,7 @@ def _init_workflow(
     config_name: Optional[str],
     temp_wf_type: Optional[str],
     queue: Optional[str] = None,
+    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
 ) -> WorkflowStatusInternal:
     wfid = (
         ctx.workflow_id
@@ -167,11 +159,13 @@ def _init_workflow(
         # Synchronously record the status and inputs for workflows and single-step workflows
         # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
         # TODO: Make this transactional (and with the queue step below)
-        dbos._sys_db.update_workflow_status(status, False, ctx.in_recovery)
-        dbos._sys_db.update_workflow_inputs(wfid, utils.serialize(inputs))
+        dbos._sys_db.update_workflow_status(
+            status, False, ctx.in_recovery, max_recovery_attempts=max_recovery_attempts
+        )
+        dbos._sys_db.update_workflow_inputs(wfid, utils.serialize_args(inputs))
     else:
         # Buffer the inputs for single-transaction workflows, but don't buffer the status
-        dbos._sys_db.buffer_workflow_inputs(wfid, utils.serialize(inputs))
+        dbos._sys_db.buffer_workflow_inputs(wfid, utils.serialize_args(inputs))
 
     if queue is not None:
         dbos._sys_db.enqueue(wfid, queue)
@@ -191,7 +185,8 @@ def _execute_workflow(
         status["status"] = "SUCCESS"
         status["output"] = utils.serialize(output)
         if status["queue_name"] is not None:
-            dbos._sys_db.remove_from_queue(status["workflow_uuid"])
+            queue = dbos._registry.queue_info_map[status["queue_name"]]
+            dbos._sys_db.remove_from_queue(status["workflow_uuid"], queue)
         dbos._sys_db.buffer_workflow_status(status)
     except DBOSWorkflowConflictIDError:
         # Retrieve the workflow handle and wait for the result.
@@ -203,9 +198,10 @@ def _execute_workflow(
         return output
     except Exception as error:
         status["status"] = "ERROR"
-        status["error"] = utils.serialize(error)
+        status["error"] = utils.serialize_exception(error)
         if status["queue_name"] is not None:
-            dbos._sys_db.remove_from_queue(status["workflow_uuid"])
+            queue = dbos._registry.queue_info_map[status["queue_name"]]
+            dbos._sys_db.remove_from_queue(status["workflow_uuid"], queue)
         dbos._sys_db.update_workflow_status(status)
         raise
 
@@ -228,7 +224,7 @@ def _execute_workflow_wthread(
         with EnterDBOSWorkflow(attributes):
             try:
                 return _execute_workflow(dbos, status, func, *args, **kwargs)
-            except Exception as e:
+            except Exception:
                 dbos.logger.error(
                     f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
                 )
@@ -299,10 +295,15 @@ def _execute_workflow_id(dbos: "DBOS", workflow_id: str) -> "WorkflowHandle[Any]
                 )
 
 
-def _workflow_wrapper(dbosreg: "_DBOSRegistry", func: F) -> F:
+def _workflow_wrapper(
+    dbosreg: "_DBOSRegistry",
+    func: F,
+    max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
+) -> F:
     func.__orig_func = func  # type: ignore
 
     fi = get_or_create_func_info(func)
+    fi.max_recovery_attempts = max_recovery_attempts
 
     @wraps(func)
     def wrapper(*args: Any, **kwargs: Any) -> Any:
@@ -335,17 +336,21 @@ def _workflow_wrapper(dbosreg: "_DBOSRegistry", func: F) -> F:
                 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),
+                max_recovery_attempts=max_recovery_attempts,
             )
 
+            dbos.logger.debug(
+                f"Running workflow, id: {ctx.workflow_id}, name: {get_dbos_func_name(func)}"
+            )
             return _execute_workflow(dbos, status, func, *args, **kwargs)
 
     wrapped_func = cast(F, wrapper)
     return wrapped_func
 
 
-def _workflow(reg: "_DBOSRegistry") -> Callable[[F], F]:
+def _workflow(reg: "_DBOSRegistry", max_recovery_attempts: int) -> Callable[[F], F]:
     def _workflow_decorator(func: F) -> F:
-        wrapped_func = _workflow_wrapper(reg, func)
+        wrapped_func = _workflow_wrapper(reg, func, max_recovery_attempts)
         reg.register_wf_function(func.__qualname__, wrapped_func)
         return wrapped_func
 
@@ -411,6 +416,7 @@ def _start_workflow(
         config_name=get_config_name(fi, func, gin_args),
         temp_wf_type=get_temp_workflow_type(func),
         queue=queue_name,
+        max_recovery_attempts=fi.max_recovery_attempts,
     )
 
     if not execute_workflow:
@@ -487,9 +493,14 @@ def _transaction(
                                     )
                                 )
                                 if recorded_output:
+                                    dbos.logger.debug(
+                                        f"Replaying transaction, id: {ctx.function_id}, name: {attributes['name']}"
+                                    )
                                     if recorded_output["error"]:
-                                        deserialized_error = utils.deserialize(
-                                            recorded_output["error"]
+                                        deserialized_error = (
+                                            utils.deserialize_exception(
+                                                recorded_output["error"]
+                                            )
                                         )
                                         has_recorded_error = True
                                         raise deserialized_error
@@ -501,6 +512,11 @@ def _transaction(
                                         raise Exception(
                                             "Output and error are both None"
                                         )
+                                else:
+                                    dbos.logger.debug(
+                                        f"Running transaction, id: {ctx.function_id}, name: {attributes['name']}"
+                                    )
+
                                 output = func(*args, **kwargs)
                                 txn_output["output"] = utils.serialize(output)
                                 assert (
@@ -527,7 +543,7 @@ def _transaction(
                         except Exception as error:
                             # Don't record the error if it was already recorded
                             if not has_recorded_error:
-                                txn_output["error"] = utils.serialize(error)
+                                txn_output["error"] = utils.serialize_exception(error)
                                 dbos._app_db.record_transaction_error(txn_output)
                             raise
             return output
@@ -598,13 +614,23 @@ def _step(
                     ctx.workflow_id, ctx.function_id
                 )
                 if recorded_output:
+                    dbos.logger.debug(
+                        f"Replaying step, id: {ctx.function_id}, name: {attributes['name']}"
+                    )
                     if recorded_output["error"] is not None:
-                        deserialized_error = utils.deserialize(recorded_output["error"])
+                        deserialized_error = utils.deserialize_exception(
+                            recorded_output["error"]
+                        )
                         raise deserialized_error
                     elif recorded_output["output"] is not None:
                         return utils.deserialize(recorded_output["output"])
                     else:
                         raise Exception("Output and error are both None")
+                else:
+                    dbos.logger.debug(
+                        f"Running step, id: {ctx.function_id}, name: {attributes['name']}"
+                    )
+
                 output = None
                 error = None
                 local_max_attempts = max_attempts if retries_allowed else 1
@@ -639,7 +665,7 @@ def _step(
                                 )
 
                 step_output["error"] = (
-                    utils.serialize(error) if error is not None else None
+                    utils.serialize_exception(error) if error is not None else None
                 )
                 dbos._sys_db.record_operation_result(step_output)
 

dbos/dbos.py

@@ -5,6 +5,7 @@ import json
 import os
 import sys
 import threading
+import traceback
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from logging import Logger
@@ -42,6 +43,7 @@ from dbos.decorators import classproperty
 from dbos.queue import Queue, queue_thread
 from dbos.recovery import _recover_pending_workflows, _startup_recovery_thread
 from dbos.registrations import (
+    DEFAULT_MAX_RECOVERY_ATTEMPTS,
     DBOSClassInfo,
     get_or_create_class_info,
     set_dbos_func_name,
@@ -338,43 +340,47 @@ class DBOS:
             _dbos_global_instance._launch()
 
     def _launch(self) -> None:
-        if self._launched:
-            dbos_logger.warning(f"DBOS was already launched")
-            return
-        self._launched = True
-        self._executor_field = ThreadPoolExecutor(max_workers=64)
-        self._sys_db_field = SystemDatabase(self.config)
-        self._app_db_field = ApplicationDatabase(self.config)
-        self._admin_server_field = AdminServer(dbos=self)
-
-        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)
-
-        # Listen to notifications
-        self._executor.submit(self._sys_db._notification_listener)
-
-        # Start flush workflow buffers thread
-        self._executor.submit(self._sys_db.flush_workflow_buffers)
-
-        # Start the queue thread
-        evt = threading.Event()
-        self.stop_events.append(evt)
-        self._executor.submit(queue_thread, evt, self)
-
-        # Grab any pollers that were deferred and start them
-        for evt, func, args, kwargs in self._registry.pollers:
+        try:
+            if self._launched:
+                dbos_logger.warning(f"DBOS was already launched")
+                return
+            self._launched = True
+            self._executor_field = ThreadPoolExecutor(max_workers=64)
+            self._sys_db_field = SystemDatabase(self.config)
+            self._app_db_field = ApplicationDatabase(self.config)
+            self._admin_server_field = AdminServer(dbos=self)
+
+            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)
+
+            # Listen to notifications
+            self._executor.submit(self._sys_db._notification_listener)
+
+            # Start flush workflow buffers thread
+            self._executor.submit(self._sys_db.flush_workflow_buffers)
+
+            # Start the queue thread
+            evt = threading.Event()
             self.stop_events.append(evt)
-            self._executor.submit(func, *args, **kwargs)
-        self._registry.pollers = []
+            self._executor.submit(queue_thread, evt, self)
 
-        dbos_logger.info("DBOS launched")
+            # Grab any pollers that were deferred and start them
+            for evt, func, args, kwargs in self._registry.pollers:
+                self.stop_events.append(evt)
+                self._executor.submit(func, *args, **kwargs)
+            self._registry.pollers = []
 
-        # 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()
+            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()
+        except Exception:
+            dbos_logger.error(f"DBOS failed to launch: {traceback.format_exc()}")
+            raise
 
     def _destroy(self) -> None:
         self._initialized = False
@@ -401,9 +407,11 @@ class DBOS:
 
     # Decorators for DBOS functionality
     @classmethod
-    def workflow(cls) -> Callable[[F], F]:
+    def workflow(
+        cls, *, max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS
+    ) -> Callable[[F], F]:
         """Decorate a function for use as a DBOS workflow."""
-        return _workflow(_get_or_create_dbos_registry())
+        return _workflow(_get_or_create_dbos_registry(), max_recovery_attempts)
 
     @classmethod
     def transaction(

dbos/error.py

@@ -32,6 +32,7 @@ class DBOSErrorCode(Enum):
     InitializationError = 3
     WorkflowFunctionNotFound = 4
     NonExistentWorkflowError = 5
+    DeadLetterQueueError = 6
     MaxStepRetriesExceeded = 7
     NotAuthorized = 8
 
@@ -86,6 +87,16 @@ class DBOSNonExistentWorkflowError(DBOSException):
         )
 
 
+class DBOSDeadLetterQueueError(DBOSException):
+    """Exception raised when a workflow database record does not exist for a given ID."""
+
+    def __init__(self, wf_id: str, max_retries: int):
+        super().__init__(
+            f"Workflow {wf_id} has been moved to the dead-letter queue after exceeding the maximum of ${max_retries} retries",
+            dbos_error_code=DBOSErrorCode.DeadLetterQueueError.value,
+        )
+
+
 class DBOSNotAuthorizedError(DBOSException):
     """Exception raised by DBOS role-based security when the user is not authorized to access a function."""
 

dbos/migrations/versions/50f3227f0b4b_fix_job_queue.py

@@ -0,0 +1,34 @@
+"""fix_job_queue
+
+Revision ID: 50f3227f0b4b
+Revises: eab0cc1d9a14
+Create Date: 2024-09-25 14:03:53.308068
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "50f3227f0b4b"
+down_revision: Union[str, None] = "eab0cc1d9a14"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.drop_constraint("job_queue_pkey", "job_queue", schema="dbos", type_="primary")
+
+    op.create_primary_key(
+        "job_queue_pkey", "job_queue", ["workflow_uuid"], schema="dbos"
+    )
+
+
+def downgrade() -> None:
+    # Reverting the changes
+    op.drop_constraint("job_queue_pkey", "job_queue", schema="dbos", type_="primary")
+
+    op.create_primary_key(
+        "job_queue_pkey", "job_queue", ["created_at_epoch_ms"], schema="dbos"
+    )

dbos/migrations/versions/d76646551a6b_job_queue_limiter.py

@@ -0,0 +1,42 @@
+"""job_queue_limiter
+
+Revision ID: d76646551a6b
+Revises: 50f3227f0b4b
+Create Date: 2024-09-25 14:48:10.218015
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "d76646551a6b"
+down_revision: Union[str, None] = "50f3227f0b4b"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "job_queue",
+        sa.Column(
+            "started_at_epoch_ms",
+            sa.BigInteger(),
+        ),
+        schema="dbos",
+    )
+    op.add_column(
+        "job_queue",
+        sa.Column(
+            "completed_at_epoch_ms",
+            sa.BigInteger(),
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("job_queue", "started_at_epoch_ms", schema="dbos")
+    op.drop_column("job_queue", "completed_at_epoch_ms", schema="dbos")

dbos/queue.py

@@ -1,18 +1,32 @@
 import threading
-import time
-from typing import TYPE_CHECKING, Optional
+import traceback
+from typing import TYPE_CHECKING, Optional, TypedDict
 
 from dbos.core import P, R, _execute_workflow_id, _start_workflow
-from dbos.error import DBOSInitializationError
 
 if TYPE_CHECKING:
     from dbos.dbos import DBOS, Workflow, WorkflowHandle
 
 
+# Limit the maximum number of functions from this queue
+# that can be started in a given period. If the limit is 5
+# and the period is 10, no more than 5 functions can be
+# started per 10 seconds.
+class Limiter(TypedDict):
+    limit: int
+    period: float
+
+
 class Queue:
-    def __init__(self, name: str, concurrency: Optional[int] = None) -> None:
+    def __init__(
+        self,
+        name: str,
+        concurrency: Optional[int] = None,
+        limiter: Optional[Limiter] = None,
+    ) -> None:
         self.name = name
         self.concurrency = concurrency
+        self.limiter = limiter
         from dbos.dbos import _get_or_create_dbos_registry
 
         registry = _get_or_create_dbos_registry()
@@ -29,8 +43,14 @@ class Queue:
 
 def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
     while not stop_event.is_set():
-        time.sleep(1)
-        for queue_name, queue in dbos._registry.queue_info_map.items():
-            wf_ids = dbos._sys_db.start_queued_workflows(queue_name, queue.concurrency)
-            for id in wf_ids:
-                _execute_workflow_id(dbos, id)
+        if stop_event.wait(timeout=1):
+            return
+        for _, queue in dbos._registry.queue_info_map.items():
+            try:
+                wf_ids = dbos._sys_db.start_queued_workflows(queue)
+                for id in wf_ids:
+                    _execute_workflow_id(dbos, id)
+            except Exception:
+                dbos.logger.warning(
+                    f"Exception encountered in queue thread: {traceback.format_exc()}"
+                )

dbos/registrations.py

@@ -3,6 +3,8 @@ from enum import Enum
 from types import FunctionType
 from typing import Any, Callable, List, Literal, Optional, Tuple, Type, cast
 
+DEFAULT_MAX_RECOVERY_ATTEMPTS = 50
+
 
 def get_dbos_func_name(f: Any) -> str:
     if hasattr(f, "dbos_function_name"):
@@ -47,6 +49,7 @@ class DBOSFuncInfo:
         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
 
 
 def get_or_create_class_info(cls: Type[Any]) -> DBOSClassInfo:

dbos/schemas/system_database.py

@@ -161,4 +161,12 @@ class SystemSchema:
             nullable=False,
             server_default=text("(EXTRACT(epoch FROM now()) * 1000::numeric)::bigint"),
         ),
+        Column(
+            "started_at_epoch_ms",
+            BigInteger(),
+        ),
+        Column(
+            "completed_at_epoch_ms",
+            BigInteger(),
+        ),
     )

dbos/system_database.py

@@ -3,7 +3,19 @@ import os
 import threading
 import time
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Sequence, Set, TypedDict, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    TypedDict,
+    cast,
+)
 
 import psycopg
 import sqlalchemy as sa
@@ -13,12 +25,20 @@ from alembic.config import Config
 from sqlalchemy.exc import DBAPIError
 
 import dbos.utils as utils
-from dbos.error import DBOSNonExistentWorkflowError, DBOSWorkflowConflictIDError
+from dbos.error import (
+    DBOSDeadLetterQueueError,
+    DBOSNonExistentWorkflowError,
+    DBOSWorkflowConflictIDError,
+)
+from dbos.registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
 
 from .dbos_config import ConfigFile
 from .logger import dbos_logger
 from .schemas.system_database import SystemSchema
 
+if TYPE_CHECKING:
+    from .queue import Queue
+
 
 class WorkflowStatusString(Enum):
     """Enumeration of values allowed for `WorkflowSatusInternal.status`."""
@@ -36,11 +56,6 @@ WorkflowStatuses = Literal[
 ]
 
 
-class WorkflowInputs(TypedDict):
-    args: Any
-    kwargs: Any
-
-
 class WorkflowStatusInternal(TypedDict):
     workflow_uuid: str
     status: WorkflowStatuses
@@ -127,7 +142,7 @@ class WorkflowInformation(TypedDict, total=False):
     # The role used to run this workflow.  Empty string if authorization is not required.
     authenticated_roles: List[str]
     # All roles the authenticated user has, if any.
-    input: Optional[WorkflowInputs]
+    input: Optional[utils.WorkflowInputs]
     output: Optional[str]
     error: Optional[str]
     request: Optional[str]
@@ -227,7 +242,9 @@ class SystemDatabase:
         status: WorkflowStatusInternal,
         replace: bool = True,
         in_recovery: bool = False,
+        *,
         conn: Optional[sa.Connection] = None,
+        max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
     ) -> None:
         cmd = pg.insert(SystemSchema.workflow_status).values(
             workflow_uuid=status["workflow_uuid"],
@@ -265,12 +282,36 @@ class SystemDatabase:
             )
         else:
             cmd = cmd.on_conflict_do_nothing()
+        cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts)  # type: ignore
 
         if conn is not None:
-            conn.execute(cmd)
+            results = conn.execute(cmd)
         else:
             with self.engine.begin() as c:
-                c.execute(cmd)
+                results = c.execute(cmd)
+        if in_recovery:
+            row = results.fetchone()
+            if row is not None:
+                recovery_attempts: int = row[0]
+                if recovery_attempts > max_recovery_attempts:
+                    with self.engine.begin() as c:
+                        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,
+                            )
+                        )
+                    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:
@@ -447,7 +488,7 @@ class SystemDatabase:
         if status == str(WorkflowStatusString.SUCCESS.value):
             return utils.deserialize(stat["output"])
         elif status == str(WorkflowStatusString.ERROR.value):
-            raise utils.deserialize(stat["error"])
+            raise utils.deserialize_exception(stat["error"])
         return None
 
     def get_workflow_info(
@@ -487,7 +528,7 @@ class SystemDatabase:
             self._exported_temp_txn_wf_status.discard(workflow_uuid)
             self._temp_txn_wf_ids.discard(workflow_uuid)
 
-    def get_workflow_inputs(self, workflow_uuid: str) -> Optional[WorkflowInputs]:
+    def get_workflow_inputs(self, workflow_uuid: str) -> Optional[utils.WorkflowInputs]:
         with self.engine.begin() as c:
             row = c.execute(
                 sa.select(SystemSchema.workflow_inputs.c.inputs).where(
@@ -496,7 +537,7 @@ class SystemDatabase:
             ).fetchone()
             if row is None:
                 return None
-            inputs: WorkflowInputs = utils.deserialize(row[0])
+            inputs: utils.WorkflowInputs = utils.deserialize_args(row[0])
             return inputs
 
     def get_workflows(self, input: GetWorkflowsInput) -> GetWorkflowsOutput:
@@ -611,7 +652,14 @@ class SystemDatabase:
                 workflow_uuid, function_id, conn=c
             )
             if recorded_output is not None:
+                dbos_logger.debug(
+                    f"Replaying send, id: {function_id}, destination_uuid: {destination_uuid}, topic: {topic}"
+                )
                 return  # Already sent before
+            else:
+                dbos_logger.debug(
+                    f"Running send, id: {function_id}, destination_uuid: {destination_uuid}, topic: {topic}"
+                )
 
             try:
                 c.execute(
@@ -647,10 +695,13 @@ class SystemDatabase:
         # First, check for previous executions.
         recorded_output = self.check_operation_execution(workflow_uuid, function_id)
         if recorded_output is not None:
+            dbos_logger.debug(f"Replaying recv, id: {function_id}, topic: {topic}")
             if recorded_output["output"] is not None:
                 return utils.deserialize(recorded_output["output"])
             else:
                 raise Exception("No output recorded in the last recv")
+        else:
+            dbos_logger.debug(f"Running recv, id: {function_id}, topic: {topic}")
 
         # Insert a condition to the notifications map, so the listener can notify it when a message is received.
         payload = f"{workflow_uuid}::{topic}"
@@ -793,9 +844,11 @@ class SystemDatabase:
         recorded_output = self.check_operation_execution(workflow_uuid, function_id)
         end_time: float
         if recorded_output is not None:
+            dbos_logger.debug(f"Replaying sleep, id: {function_id}, seconds: {seconds}")
             assert recorded_output["output"] is not None, "no recorded end time"
             end_time = utils.deserialize(recorded_output["output"])
         else:
+            dbos_logger.debug(f"Running sleep, id: {function_id}, seconds: {seconds}")
             end_time = time.time() + seconds
             try:
                 self.record_operation_result(
@@ -825,7 +878,10 @@ class SystemDatabase:
                 workflow_uuid, function_id, conn=c
             )
             if recorded_output is not None:
+                dbos_logger.debug(f"Replaying set_event, id: {function_id}, key: {key}")
                 return  # Already sent before
+            else:
+                dbos_logger.debug(f"Running set_event, id: {function_id}, key: {key}")
 
             c.execute(
                 pg.insert(SystemSchema.workflow_events)
@@ -866,10 +922,17 @@ class SystemDatabase:
                 caller_ctx["workflow_uuid"], caller_ctx["function_id"]
             )
             if recorded_output is not None:
+                dbos_logger.debug(
+                    f"Replaying get_event, id: {caller_ctx['function_id']}, key: {key}"
+                )
                 if recorded_output["output"] is not None:
                     return utils.deserialize(recorded_output["output"])
                 else:
                     raise Exception("No output recorded in the last get_event")
+            else:
+                dbos_logger.debug(
+                    f"Running get_event, id: {caller_ctx['function_id']}, key: {key}"
+                )
 
         payload = f"{target_uuid}::{key}"
         condition = threading.Condition()
@@ -1024,22 +1087,60 @@ class SystemDatabase:
                 .on_conflict_do_nothing()
             )
 
-    def start_queued_workflows(
-        self, queue_name: str, concurrency: Optional[int]
-    ) -> List[str]:
+    def start_queued_workflows(self, queue: "Queue") -> List[str]:
+        start_time_ms = int(time.time() * 1000)
+        if queue.limiter is not None:
+            limiter_period_ms = int(queue.limiter["period"] * 1000)
         with self.engine.begin() as c:
-            query = sa.select(SystemSchema.job_queue.c.workflow_uuid).where(
-                SystemSchema.job_queue.c.queue_name == queue_name
+            # Execute with snapshot isolation to ensure multiple workers respect limits
+            c.execute(sa.text("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"))
+
+            # If there is a limiter, compute how many functions have started in its period.
+            if queue.limiter is not None:
+                query = (
+                    sa.select(sa.func.count())
+                    .select_from(SystemSchema.job_queue)
+                    .where(SystemSchema.job_queue.c.started_at_epoch_ms.isnot(None))
+                    .where(
+                        SystemSchema.job_queue.c.started_at_epoch_ms
+                        > start_time_ms - limiter_period_ms
+                    )
+                )
+                num_recent_queries = c.execute(query).fetchone()[0]  # type: ignore
+                if num_recent_queries >= queue.limiter["limit"]:
+                    return []
+
+            # Select not-yet-completed functions in the queue ordered by the
+            # time at which they were enqueued.
+            # If there is a concurrency limit N, select only the N most recent
+            # functions, else select all of them.
+            query = (
+                sa.select(
+                    SystemSchema.job_queue.c.workflow_uuid,
+                    SystemSchema.job_queue.c.started_at_epoch_ms,
+                )
+                .where(SystemSchema.job_queue.c.queue_name == queue.name)
+                .where(SystemSchema.job_queue.c.completed_at_epoch_ms == None)
+                .order_by(SystemSchema.job_queue.c.created_at_epoch_ms.asc())
             )
-            if concurrency is not None:
-                query = query.order_by(
-                    SystemSchema.job_queue.c.created_at_epoch_ms.asc()
-                ).limit(concurrency)
+            if queue.concurrency is not None:
+                query = query.limit(queue.concurrency)
+
+            # From the functions retrieved, get the workflow IDs of the functions
+            # that have not yet been started so we can start them.
             rows = c.execute(query).fetchall()
-            dequeued_ids: List[str] = [row[0] for row in rows]
-            ret_ids = []
+            dequeued_ids: List[str] = [row[0] for row in rows if row[1] is None]
+            ret_ids: list[str] = []
             for id in dequeued_ids:
-                result = c.execute(
+
+                # 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:
+                    if len(ret_ids) + num_recent_queries >= queue.limiter["limit"]:
+                        break
+
+                # To start a function, first set its status to PENDING
+                c.execute(
                     SystemSchema.workflow_status.update()
                     .where(SystemSchema.workflow_status.c.workflow_uuid == id)
                     .where(
@@ -1048,14 +1149,42 @@ class SystemDatabase:
                     )
                     .values(status=WorkflowStatusString.PENDING.value)
                 )
-                if result.rowcount > 0:
-                    ret_ids.append(id)
+
+                # Then give it a start time
+                c.execute(
+                    SystemSchema.job_queue.update()
+                    .where(SystemSchema.job_queue.c.workflow_uuid == id)
+                    .values(started_at_epoch_ms=start_time_ms)
+                )
+                ret_ids.append(id)
+
+            # If we have a limiter, garbage-collect all completed functions started
+            # before the period. If there's no limiter, there's no need--they were
+            # deleted on completion.
+            if queue.limiter is not None:
+                c.execute(
+                    sa.delete(SystemSchema.job_queue)
+                    .where(SystemSchema.job_queue.c.completed_at_epoch_ms != None)
+                    .where(
+                        SystemSchema.job_queue.c.started_at_epoch_ms
+                        < start_time_ms - limiter_period_ms
+                    )
+                )
+
+            # Return the IDs of all functions we started
             return ret_ids
 
-    def remove_from_queue(self, workflow_id: str) -> None:
+    def remove_from_queue(self, workflow_id: str, queue: "Queue") -> None:
         with self.engine.begin() as c:
-            c.execute(
-                sa.delete(SystemSchema.job_queue).where(
-                    SystemSchema.job_queue.c.workflow_uuid == workflow_id
+            if queue.limiter is None:
+                c.execute(
+                    sa.delete(SystemSchema.job_queue).where(
+                        SystemSchema.job_queue.c.workflow_uuid == workflow_id
+                    )
+                )
+            else:
+                c.execute(
+                    sa.update(SystemSchema.job_queue)
+                    .where(SystemSchema.job_queue.c.workflow_uuid == workflow_id)
+                    .values(completed_at_epoch_ms=int(time.time() * 1000))
                 )
-            )

dbos/utils.py

@@ -1,14 +1,55 @@
-from typing import Any
+import types
+from typing import Any, Dict, Tuple, TypedDict
 
 import jsonpickle  # type: ignore
 
 
+class WorkflowInputs(TypedDict):
+    args: Tuple[Any, ...]
+    kwargs: Dict[str, Any]
+
+
+def validate_item(data: Any) -> None:
+    if isinstance(data, (types.FunctionType, types.MethodType)):
+        raise TypeError("Serialized data item should not be a function")
+
+
 def serialize(data: Any) -> str:
     """Serialize an object to a JSON string using jsonpickle."""
-    encoded_data: str = jsonpickle.encode(data)
+    validate_item(data)
+    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    return encoded_data
+
+
+def serialize_args(data: WorkflowInputs) -> str:
+    """Serialize args to a JSON string using jsonpickle."""
+    arg: Any
+    for arg in data["args"]:
+        validate_item(arg)
+    for arg in data["kwargs"].values():
+        validate_item(arg)
+    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    return encoded_data
+
+
+def serialize_exception(data: Exception) -> str:
+    """Serialize an Exception object to a JSON string using jsonpickle."""
+    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
     return encoded_data
 
 
 def deserialize(serialized_data: str) -> Any:
     """Deserialize a JSON string back to a Python object using jsonpickle."""
     return jsonpickle.decode(serialized_data)
+
+
+def deserialize_args(serialized_data: str) -> WorkflowInputs:
+    """Deserialize a JSON string back to a Python object list using jsonpickle."""
+    args: WorkflowInputs = jsonpickle.decode(serialized_data)
+    return args
+
+
+def deserialize_exception(serialized_data: str) -> Exception:
+    """Deserialize JSON string back to a Python Exception using jsonpickle."""
+    upo: Exception = jsonpickle.decode(serialized_data)
+    return upo

{dbos-0.8.0a0.dist-info → dbos-0.8.0a7.dist-info}/METADATA

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

{dbos-0.8.0a0.dist-info → dbos-0.8.0a7.dist-info}/RECORD

@@ -1,18 +1,18 @@
-dbos-0.8.0a0.dist-info/METADATA,sha256=25gzYG9xJsemmuA1r8k_jnJyp3z3hHnRoc2UzAR2CQI,5010
-dbos-0.8.0a0.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
-dbos-0.8.0a0.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
-dbos-0.8.0a0.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
-dbos/__init__.py,sha256=jjlBFzSAzO2e-LD5IKJw7bRqZjrxpF5Sn-_JUJJptHU,680
+dbos-0.8.0a7.dist-info/METADATA,sha256=cqNzAf7TbC3n_DDGTPUj_xawggwMydLnIQ5056iFDUI,5010
+dbos-0.8.0a7.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
+dbos-0.8.0a7.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
+dbos-0.8.0a7.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos/__init__.py,sha256=-h1QgWNL11CiLlHEKa2ycAJVJw5SXYZ4BGNNWBAiE9k,726
 dbos/admin_sever.py,sha256=Qg5T3YRrbPW05PR_99yAaxgo1ugQrAp_uTeTqSfjm_k,3397
 dbos/application_database.py,sha256=knFK8We8y6WrIpnFCKvFq5hvSuFQqUuJqOqDpSVMCPI,5521
 dbos/cli.py,sha256=z5dXbbnGWzSC3E1rfS8Lp1_OIImzcDKM7jP-iu_Q4aI,8602
-dbos/context.py,sha256=1Xp0i-ZP72Vj6eMdHuyfXi9RHnoT2w4MZ2Kx1CrKnQ8,16559
-dbos/core.py,sha256=LnQRHLjfpdlzlRkElduay8nObNMyRNC8SJSZEzqei7o,29161
+dbos/context.py,sha256=4MsxZdoh1WIsgoUsaxo0B6caGN6xq2WC60MzbBppzGk,17738
+dbos/core.py,sha256=ggsRC2XicvNI1qqruEFoqxoTU5oSSnhMZvDih3AG_3A,30879
 dbos/dbos-config.schema.json,sha256=azpfmoDZg7WfSy3kvIsk9iEiKB_-VZt03VEOoXJAkqE,5331
-dbos/dbos.py,sha256=RtDcvKe4sm1TlnCGU4cyex-UI7hMMlhgzmOl1NuRLo4,29294
+dbos/dbos.py,sha256=LnqX7rFETpcyxT9YHs4Uc3uOB4EDQC-zis3UFQU4smc,29705
 dbos/dbos_config.py,sha256=NJVze2GkKgYUmcPP31Unb-QpsA0TzImEeQGJgVq6W6k,5352
 dbos/decorators.py,sha256=lbPefsLK6Cya4cb7TrOcLglOpGT3pc6qjZdsQKlfZLg,629
-dbos/error.py,sha256=c2y7d3Cbb-ZOO-M9txcCxuyPE1bdnwKrJnXgJXYb-pQ,3437
+dbos/error.py,sha256=UETk8CoZL-TO2Utn1-E7OSWelhShWmKM-fOlODMR9PE,3893
 dbos/fastapi.py,sha256=gx9hlpxYOiwbuhSlbY9bn5C-F_FsCbrJvkX9ZAvDG6U,3418
 dbos/flask.py,sha256=azr4geMEGuuTBCyxIZmgDmmP-6s_pTIF-lGyp9Q4IB8,2430
 dbos/kafka.py,sha256=LH3hbNapnkjLcuXNUtdGU0398JafWb-t0GwUl3LOzkc,3645
@@ -20,21 +20,23 @@ dbos/kafka_message.py,sha256=NYvOXNG3Qn7bghn1pv3fg4Pbs86ILZGcK4IB-MLUNu0,409
 dbos/logger.py,sha256=D-aFSZUCHBP34J1IZ5YNkTrJW-rDiH3py_v9jLU4Yrk,3565
 dbos/migrations/env.py,sha256=38SIGVbmn_VV2x2u1aHLcPOoWgZ84eCymf3g_NljmbU,1626
 dbos/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
+dbos/migrations/versions/50f3227f0b4b_fix_job_queue.py,sha256=ZtnsZFMuon-D0n8V5BR10jQEqJPUsYsOwt29FAoKG8g,868
 dbos/migrations/versions/5c361fc04708_added_system_tables.py,sha256=QMgFMb0aLgC25YicsvPSr6AHRCA6Zd66hyaRUhwKzrQ,6404
 dbos/migrations/versions/a3b18ad34abe_added_triggers.py,sha256=Rv0ZsZYZ_WdgGEULYsPfnp4YzaO5L198gDTgYY39AVA,2022
+dbos/migrations/versions/d76646551a6b_job_queue_limiter.py,sha256=M1upulBOLXm9ORJc2Q6PA0AwX0CA3zgqc-NhYS-eNPY,948
 dbos/migrations/versions/eab0cc1d9a14_job_queue.py,sha256=_9-FCW-zOpCQfblTS_yRLtFiUaWlC1tM4BoKBTDeH9k,1395
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
-dbos/queue.py,sha256=5unMPX1ThoVWfXOR2IUSeIjgATgcX_YmZS3Fz1Uoc7o,1228
+dbos/queue.py,sha256=4MIWYDdl3DhYsnayy747xF1Jgdq2qvV4On3KJbJ4NDU,1764
 dbos/recovery.py,sha256=zqtO_ExGoIErLMVnbneU3VeHLVWvhV4jnfqssAVlQQk,2016
-dbos/registrations.py,sha256=gMI-u05tv5bpvyddQGtoUgCsqARx51aOY7p0JXPafQo,6539
+dbos/registrations.py,sha256=mei6q6_3R5uei8i_Wo_TqGZs85s10shOekDX41sFYD0,6642
 dbos/request.py,sha256=-FIwtknayvRl6OjvqO4V2GySVzSdP1Ft3cc9ZBS-PLY,928
 dbos/roles.py,sha256=7Lh7uwUq1dpa6TXCOHre4mPTd5qmXzK_QPkvYR52DXg,2285
 dbos/scheduler/croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
 dbos/scheduler/scheduler.py,sha256=Sz4EIpAtur7so2YajTic64GrTpa4qPw8QxXn0M34v80,1360
 dbos/schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/schemas/application_database.py,sha256=q_Wr2XbiZNBYFkOtu7uKavo1T_cSOBblxKGHThYGGsY,962
-dbos/schemas/system_database.py,sha256=ed4c1UntsD-cqXD0ekM4jvcYYEViavDh_G6c0pVDe7k,4938
-dbos/system_database.py,sha256=VSGFSPubbMokGYsZfRb6cQPltLfjoWM-Re_2Gj9qkRc,41844
+dbos/schemas/system_database.py,sha256=-37sNXfx6cNGyzndj9mrWQLDH5iIBrwsT56ZiJ56Sj0,5116
+dbos/system_database.py,sha256=x-TpKHIS187KED6BxPozvMLt6Qjhgh3dYSKZkd6epM0,47764
 dbos/templates/hello/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/templates/hello/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/templates/hello/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
@@ -46,6 +48,6 @@ dbos/templates/hello/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy
 dbos/templates/hello/migrations/versions/2024_07_31_180642_init.py,sha256=U5thFWGqNN4QLrNXT7wUUqftIFDNE5eSdqD8JNW1mec,942
 dbos/templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKse9HsWv5fynJEFJHM,1681
 dbos/tracer.py,sha256=GaXDhdKKF_IQp5SAMipGXiDVwteRKjNbrXyYCH1mor0,2520
-dbos/utils.py,sha256=hWj9iWDrby2cVEhb0pG-IdnrxLqP64NhkaWUXiLc8bA,402
+dbos/utils.py,sha256=lwRymY-y7GprAS8pKmbICQvOJd5eGxKGTxCMFn0OwaQ,1739
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.8.0a0.dist-info/RECORD,,
+dbos-0.8.0a7.dist-info/RECORD,,