dbos 0.5.0a3__py3-none-any.whl → 0.5.0a5__py3-none-any.whl

dbos/application_database.py

@@ -3,6 +3,7 @@ from typing import Optional, TypedDict, cast
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 import sqlalchemy.exc as sa_exc
+from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
 from dbos.error import DBOSWorkflowConflictIDError
@@ -95,8 +96,10 @@ class ApplicationDatabase:
                     ),
                 )
             )
-        except sa_exc.IntegrityError:
-            raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+            raise dbapi_error
         except Exception as e:
             raise e
 
@@ -118,8 +121,10 @@ class ApplicationDatabase:
                         ),
                     )
                 )
-        except sa_exc.IntegrityError:
-            raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+            raise dbapi_error
         except Exception as e:
             raise e
 

dbos/context.py

@@ -18,18 +18,16 @@ from .logger import dbos_logger
 from .tracer import dbos_tracer
 
 
-# Values must be the same as in TypeScript Transact
+# These are used to tag OTel traces
 class OperationType(Enum):
     HANDLER = "handler"
     WORKFLOW = "workflow"
     TRANSACTION = "transaction"
-    COMMUNICATOR = "communicator"
+    STEP = "step"
     PROCEDURE = "procedure"
 
 
-OperationTypes = Literal[
-    "handler", "workflow", "transaction", "communicator", "procedure"
-]
+OperationTypes = Literal["handler", "workflow", "transaction", "step", "procedure"]
 
 
 # Keys must be the same as in TypeScript Transact
@@ -64,7 +62,7 @@ class DBOSContext:
         self.function_id: int = -1
         self.in_recovery: bool = False
 
-        self.curr_comm_function_id: int = -1
+        self.curr_step_function_id: int = -1
         self.curr_tx_function_id: int = -1
         self.sql_session: Optional[Session] = None
         self.spans: list[Span] = []
@@ -117,26 +115,26 @@ class DBOSContext:
     def is_workflow(self) -> bool:
         return (
             len(self.workflow_id) > 0
-            and not self.is_communicator()
+            and not self.is_step()
             and not self.is_transaction()
         )
 
     def is_transaction(self) -> bool:
         return self.sql_session is not None
 
-    def is_communicator(self) -> bool:
-        return self.curr_comm_function_id >= 0
+    def is_step(self) -> bool:
+        return self.curr_step_function_id >= 0
 
-    def start_communicator(
+    def start_step(
         self,
         fid: int,
         attributes: TracedAttributes,
     ) -> None:
-        self.curr_comm_function_id = fid
+        self.curr_step_function_id = fid
         self._start_span(attributes)
 
-    def end_communicator(self, exc_value: Optional[BaseException]) -> None:
-        self.curr_comm_function_id = -1
+    def end_step(self, exc_value: Optional[BaseException]) -> None:
+        self.curr_step_function_id = -1
         self._end_span(exc_value)
 
     def start_transaction(
@@ -375,7 +373,7 @@ class EnterDBOSChildWorkflow:
     def __enter__(self) -> DBOSContext:
         ctx = assert_current_dbos_context()
         self.parent_ctx = ctx
-        assert ctx.is_workflow()  # Is in a workflow and not in tx/comm
+        assert ctx.is_workflow()  # Is in a workflow and not in a step
         ctx.function_id += 1
         if len(ctx.id_assigned_for_next_workflow) == 0:
             ctx.id_assigned_for_next_workflow = (
@@ -401,7 +399,7 @@ class EnterDBOSChildWorkflow:
         return False  # Did not handle
 
 
-class EnterDBOSCommunicator:
+class EnterDBOSStep:
     def __init__(
         self,
         attributes: TracedAttributes,
@@ -412,7 +410,7 @@ class EnterDBOSCommunicator:
         ctx = assert_current_dbos_context()
         assert ctx.is_workflow()
         ctx.function_id += 1
-        ctx.start_communicator(ctx.function_id, attributes=self.attributes)
+        ctx.start_step(ctx.function_id, attributes=self.attributes)
         return ctx
 
     def __exit__(
@@ -422,8 +420,8 @@ class EnterDBOSCommunicator:
         traceback: Optional[TracebackType],
     ) -> Literal[False]:
         ctx = assert_current_dbos_context()
-        assert ctx.is_communicator()
-        ctx.end_communicator(exc_value)
+        assert ctx.is_step()
+        ctx.end_step(exc_value)
         return False  # Did not handle
 
 

dbos/core.py

@@ -19,20 +19,18 @@ from dbos.context import (
     DBOSContextEnsure,
     DBOSContextSwap,
     EnterDBOSChildWorkflow,
-    EnterDBOSCommunicator,
+    EnterDBOSStep,
     EnterDBOSTransaction,
     EnterDBOSWorkflow,
     OperationType,
     SetWorkflowID,
     TracedAttributes,
     assert_current_dbos_context,
-    clear_local_dbos_context,
     get_local_dbos_context,
-    set_local_dbos_context,
 )
 from dbos.error import (
-    DBOSCommunicatorMaxRetriesExceededError,
     DBOSException,
+    DBOSMaxStepRetriesExceeded,
     DBOSNonExistentWorkflowError,
     DBOSRecoveryError,
     DBOSWorkflowConflictIDError,
@@ -61,7 +59,6 @@ if TYPE_CHECKING:
     from dbos.dbos import IsolationLevel
 
 from sqlalchemy.exc import DBAPIError
-from sqlalchemy.orm import Session
 
 P = ParamSpec("P")  # A generic type for workflow parameters
 R = TypeVar("R", covariant=True)  # A generic type for workflow return values
@@ -139,14 +136,18 @@ def _init_workflow(
         "recovery_attempts": None,
     }
 
+    # If we have a class name, the first arg is the instance and do not serialize
+    if class_name is not None:
+        inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
+
     if temp_wf_type != "transaction":
-        # Don't synchronously record the status or inputs for single transaction workflows
+        # 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
         dbos.sys_db.update_workflow_status(status, False, ctx.in_recovery)
-
-        # If we have a class name, the first arg is the instance and do not serialize
-        if class_name is not None:
-            inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
         dbos.sys_db.update_workflow_inputs(wfid, utils.serialize(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))
 
     return status
 
@@ -176,20 +177,6 @@ def _execute_workflow(
         status["error"] = utils.serialize(error)
         dbos.sys_db.update_workflow_status(status)
         raise error
-    finally:
-        if get_temp_workflow_type(func) == "transaction":
-            # Buffer the inputs for single transaction workflows
-            inputs: WorkflowInputs = {
-                "args": args,
-                "kwargs": kwargs,
-            }
-            # If we have a class name, the first arg is the instance and do not serialize
-            class_name = get_dbos_class_name(get_func_info(func), func, args)
-            if class_name is not None:
-                inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
-            dbos.sys_db.buffer_workflow_inputs(
-                status["workflow_uuid"], utils.serialize(inputs)
-            )
 
     return output
 
@@ -362,7 +349,7 @@ def _start_workflow(
 
     # Sequence of events for starting a workflow:
     #   First - is there a WF already running?
-    #      (and not in tx/comm as that is an error)
+    #      (and not in step as that is an error)
     #   Assign an ID to the workflow, if it doesn't have an app-assigned one
     #      If this is a root workflow, assign a new ID
     #      If this is a child workflow, assign parent wf id with call# suffix
@@ -370,7 +357,7 @@ def _start_workflow(
     #   Pass the new context to a worker thread that will run the wf function
     cur_ctx = get_local_dbos_context()
     if cur_ctx is not None and cur_ctx.is_within_workflow():
-        assert cur_ctx.is_workflow()  # Not in tx / comm
+        assert cur_ctx.is_workflow()  # Not in a step
         cur_ctx.function_id += 1
         if len(cur_ctx.id_assigned_for_next_workflow) == 0:
             cur_ctx.id_assigned_for_next_workflow = (
@@ -517,7 +504,7 @@ def _transaction(
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             rr: Optional[str] = check_required_roles(func, fi)
             # Entering transaction is allowed:
-            #  In a workflow (that is not in a transaction / comm already)
+            #  In a workflow (that is not in a step already)
             #  Not in a workflow (we will start the single op workflow)
             ctx = get_local_dbos_context()
             if ctx and ctx.is_within_workflow():
@@ -544,7 +531,7 @@ def _transaction(
     return decorator
 
 
-def _communicator(
+def _step(
     dbosreg: "_DBOSRegistry",
     *,
     retries_allowed: bool = False,
@@ -554,7 +541,7 @@ def _communicator(
 ) -> Callable[[F], F]:
     def decorator(func: F) -> F:
 
-        def invoke_comm(*args: Any, **kwargs: Any) -> Any:
+        def invoke_step(*args: Any, **kwargs: Any) -> Any:
             if dbosreg.dbos is None:
                 raise DBOSException(
                     f"Function {func.__name__} invoked before DBOS initialized"
@@ -563,10 +550,10 @@ def _communicator(
 
             attributes: TracedAttributes = {
                 "name": func.__name__,
-                "operationType": OperationType.COMMUNICATOR.value,
+                "operationType": OperationType.STEP.value,
             }
-            with EnterDBOSCommunicator(attributes) as ctx:
-                comm_output: OperationResultInternal = {
+            with EnterDBOSStep(attributes) as ctx:
+                step_output: OperationResultInternal = {
                     "workflow_uuid": ctx.workflow_id,
                     "function_id": ctx.function_id,
                     "output": None,
@@ -591,24 +578,24 @@ def _communicator(
                 for attempt in range(1, local_max_attempts + 1):
                     try:
                         output = func(*args, **kwargs)
-                        comm_output["output"] = utils.serialize(output)
+                        step_output["output"] = utils.serialize(output)
                         error = None
                         break
                     except Exception as err:
                         error = err
                         if retries_allowed:
                             dbos.logger.warning(
-                                f"Communicator being automatically retried. (attempt {attempt} of {local_max_attempts}). {traceback.format_exc()}"
+                                f"Step being automatically retried. (attempt {attempt} of {local_max_attempts}). {traceback.format_exc()}"
                             )
                             ctx.get_current_span().add_event(
-                                f"Communicator attempt {attempt} failed",
+                                f"Step attempt {attempt} failed",
                                 {
                                     "error": str(error),
                                     "retryIntervalSeconds": local_interval_seconds,
                                 },
                             )
                             if attempt == local_max_attempts:
-                                error = DBOSCommunicatorMaxRetriesExceededError()
+                                error = DBOSMaxStepRetriesExceeded()
                             else:
                                 time.sleep(local_interval_seconds)
                                 local_interval_seconds = min(
@@ -616,10 +603,11 @@ def _communicator(
                                     max_retry_interval_seconds,
                                 )
 
-                comm_output["error"] = (
+                step_output["error"] = (
                     utils.serialize(error) if error is not None else None
                 )
-                dbos.sys_db.record_operation_result(comm_output)
+                dbos.sys_db.record_operation_result(step_output)
+
                 if error is not None:
                     raise error
                 return output
@@ -629,20 +617,18 @@ def _communicator(
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             rr: Optional[str] = check_required_roles(func, fi)
-            # Entering communicator is allowed:
-            #  In a communicator already, just call the original function directly.
-            #  In a workflow (that is not in a transaction / comm already)
+            # Entering step is allowed:
+            #  In a step already, just call the original function directly.
+            #  In a workflow (that is not in a step already)
             #  Not in a workflow (we will start the single op workflow)
             ctx = get_local_dbos_context()
-            if ctx and ctx.is_communicator():
+            if ctx and ctx.is_step():
                 # Call the original function directly
                 return func(*args, **kwargs)
             if ctx and ctx.is_within_workflow():
-                assert (
-                    ctx.is_workflow()
-                ), "Communicators must be called from within workflows"
+                assert ctx.is_workflow(), "Steps must be called from within workflows"
                 with DBOSAssumeRole(rr):
-                    return invoke_comm(*args, **kwargs)
+                    return invoke_step(*args, **kwargs)
             else:
                 tempwf = dbosreg.workflow_info_map.get("<temp>." + func.__qualname__)
                 assert tempwf
@@ -653,7 +639,7 @@ def _communicator(
 
         wrapped_wf = _workflow_wrapper(dbosreg, temp_wf)
         set_dbos_func_name(temp_wf, "<temp>." + func.__qualname__)
-        set_temp_workflow_type(temp_wf, "communicator")
+        set_temp_workflow_type(temp_wf, "step")
         dbosreg.register_wf_function(get_dbos_func_name(temp_wf), wrapped_wf)
 
         return cast(F, wrapper)
@@ -668,10 +654,10 @@ def _send(
         attributes: TracedAttributes = {
             "name": "send",
         }
-        with EnterDBOSCommunicator(attributes) as ctx:
+        with EnterDBOSStep(attributes) as ctx:
             dbos.sys_db.send(
                 ctx.workflow_id,
-                ctx.curr_comm_function_id,
+                ctx.curr_step_function_id,
                 destination_id,
                 message,
                 topic,
@@ -697,12 +683,12 @@ def _recv(
         attributes: TracedAttributes = {
             "name": "recv",
         }
-        with EnterDBOSCommunicator(attributes) as ctx:
+        with EnterDBOSStep(attributes) as ctx:
             ctx.function_id += 1  # Reserve for the sleep
             timeout_function_id = ctx.function_id
             return dbos.sys_db.recv(
                 ctx.workflow_id,
-                ctx.curr_comm_function_id,
+                ctx.curr_step_function_id,
                 timeout_function_id,
                 topic,
                 timeout_seconds,
@@ -722,9 +708,9 @@ def _set_event(dbos: "DBOS", key: str, value: Any) -> None:
         attributes: TracedAttributes = {
             "name": "set_event",
         }
-        with EnterDBOSCommunicator(attributes) as ctx:
+        with EnterDBOSStep(attributes) as ctx:
             dbos.sys_db.set_event(
-                ctx.workflow_id, ctx.curr_comm_function_id, key, value
+                ctx.workflow_id, ctx.curr_step_function_id, key, value
             )
     else:
         # Cannot call it from outside of a workflow
@@ -743,12 +729,12 @@ def _get_event(
         attributes: TracedAttributes = {
             "name": "get_event",
         }
-        with EnterDBOSCommunicator(attributes) as ctx:
+        with EnterDBOSStep(attributes) as ctx:
             ctx.function_id += 1
             timeout_function_id = ctx.function_id
             caller_ctx: GetEventWorkflowContext = {
                 "workflow_uuid": ctx.workflow_id,
-                "function_id": ctx.curr_comm_function_id,
+                "function_id": ctx.curr_step_function_id,
                 "timeout_function_id": timeout_function_id,
             }
             return dbos.sys_db.get_event(workflow_id, key, timeout_seconds, caller_ctx)

dbos/dbos.py

@@ -25,13 +25,13 @@ from opentelemetry.trace import Span
 
 from dbos.core import (
     TEMP_SEND_WF_NAME,
-    _communicator,
     _execute_workflow_id,
     _get_event,
     _recv,
     _send,
     _set_event,
     _start_workflow,
+    _step,
     _transaction,
     _workflow,
     _workflow_wrapper,
@@ -63,7 +63,7 @@ else:
 
 from dbos.admin_sever import AdminServer
 from dbos.context import (
-    EnterDBOSCommunicator,
+    EnterDBOSStep,
     TracedAttributes,
     assert_current_dbos_context,
     get_local_dbos_context,
@@ -179,7 +179,7 @@ class DBOS:
     Main access class for DBOS functionality.
 
     `DBOS` contains functions and properties for:
-    1. Decorating classes, workflows, transactions, and communicators
+    1. Decorating classes, workflows, and steps
     2. Starting workflow functions
     3. Retrieving workflow status information
     4. Interacting with workflows via events and messages
@@ -388,9 +388,8 @@ class DBOS:
         """
         return _transaction(_get_or_create_dbos_registry(), isolation_level)
 
-    # Mirror the CommunicatorConfig from TS. However, we disable retries by default.
     @classmethod
-    def communicator(
+    def step(
         cls,
         *,
         retries_allowed: bool = False,
@@ -399,17 +398,17 @@ class DBOS:
         backoff_rate: float = 2.0,
     ) -> Callable[[F], F]:
         """
-        Decorate and configure a function for use as a DBOS communicator.
+        Decorate and configure a function for use as a DBOS step.
 
         Args:
             retries_allowed(bool): If true, enable retries on thrown exceptions
             interval_seconds(float): Time between retry attempts
             backoff_rate(float): Multiplier for exponentially increasing `interval_seconds` between retries
-            max_attempts(int): Maximum number of communicator retries before raising an exception
+            max_attempts(int): Maximum number of retries before raising an exception
 
         """
 
-        return _communicator(
+        return _step(
             _get_or_create_dbos_registry(),
             retries_allowed=retries_allowed,
             interval_seconds=interval_seconds,
@@ -542,9 +541,9 @@ class DBOS:
         }
         if seconds <= 0:
             return
-        with EnterDBOSCommunicator(attributes) as ctx:
+        with EnterDBOSStep(attributes) as ctx:
             _get_dbos_instance().sys_db.sleep(
-                ctx.workflow_id, ctx.curr_comm_function_id, seconds
+                ctx.workflow_id, ctx.curr_step_function_id, seconds
             )
 
     @classmethod
@@ -615,9 +614,7 @@ class DBOS:
     def sql_session(cls) -> Session:
         """Return the SQLAlchemy `Session` for the current context, which must be within a transaction function."""
         ctx = assert_current_dbos_context()
-        assert (
-            ctx.is_transaction()
-        ), "sql_session is only available within a transaction."
+        assert ctx.is_transaction(), "db is only available within a transaction."
         rv = ctx.sql_session
         assert rv
         return rv
@@ -628,7 +625,7 @@ class DBOS:
         ctx = assert_current_dbos_context()
         assert (
             ctx.is_within_workflow()
-        ), "workflow_id is only available within a workflow, transaction, or communicator."
+        ), "workflow_id is only available within a DBOS operation."
         return ctx.workflow_id
 
     @classproperty
@@ -722,9 +719,9 @@ class DBOSConfiguredInstance:
     """
     Base class for classes containing DBOS member functions.
 
-    When a class contains workflow, transaction, and communicator functions that access
-    instance state, the DBOS workflow executor needs a name for the instance.  This name
-    is recorded in the database, and used to refer to the proper instance upon recovery.
+    When a class contains DBOS functions that access instance state, the DBOS workflow
+    executor needs a name for the instance.  This name is recorded in the database, and
+    used to refer to the proper instance upon recovery.
 
     Use `DBOSConfiguredInstance` to specify the instance name and register the instance
     with the DBOS workflow executor.

dbos/error.py

@@ -32,7 +32,7 @@ class DBOSErrorCode(Enum):
     WorkflowFunctionNotFound = 4
     NonExistentWorkflowError = 5
     DuplicateWorkflowEventError = 6
-    CommunicatorMaxRetriesExceeded = 7
+    MaxStepRetriesExceeded = 7
     NotAuthorized = 8
 
 
@@ -106,11 +106,11 @@ class DBOSNotAuthorizedError(DBOSException):
         )
 
 
-class DBOSCommunicatorMaxRetriesExceededError(DBOSException):
-    """Exception raised when a communicator function was retried the maximimum number of times without success."""
+class DBOSMaxStepRetriesExceeded(DBOSException):
+    """Exception raised when a step was retried the maximimum number of times without success."""
 
     def __init__(self) -> None:
         super().__init__(
-            "Communicator reached maximum retries.",
-            dbos_error_code=DBOSErrorCode.CommunicatorMaxRetriesExceeded.value,
+            "Step reached maximum retries.",
+            dbos_error_code=DBOSErrorCode.MaxStepRetriesExceeded.value,
         )

dbos/registrations.py

@@ -16,7 +16,7 @@ def set_dbos_func_name(f: Any, name: str) -> None:
     setattr(f, "dbos_function_name", name)
 
 
-TempWorkflowType = Literal["transaction", "communicator", "send", None]
+TempWorkflowType = Literal["transaction", "step", "send", None]
 
 
 def get_temp_workflow_type(f: Any) -> TempWorkflowType:

dbos/system_database.py

@@ -12,6 +12,7 @@ from typing import (
     Literal,
     Optional,
     Sequence,
+    Set,
     TypedDict,
     cast,
 )
@@ -21,6 +22,7 @@ import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 from alembic import command
 from alembic.config import Config
+from sqlalchemy.exc import DBAPIError
 
 import dbos.utils as utils
 from dbos.error import (
@@ -33,9 +35,6 @@ from .dbos_config import ConfigFile
 from .logger import dbos_logger
 from .schemas.system_database import SystemSchema
 
-if TYPE_CHECKING:
-    from dbos.dbos import DBOS
-
 
 class WorkflowStatusString(Enum):
     """Enumeration of values allowed for `WorkflowSatusInternal.status`."""
@@ -212,6 +211,9 @@ class SystemDatabase:
         # Initialize the workflow status and inputs buffers
         self._workflow_status_buffer: Dict[str, WorkflowStatusInternal] = {}
         self._workflow_inputs_buffer: Dict[str, str] = {}
+        # Two sets for tracking which single-transaction workflows have been exported to the status table
+        self._exported_temp_txn_wf_status: Set[str] = set()
+        self._temp_txn_wf_ids: Set[str] = set()
         self._is_flushing_status_buffer = False
 
         # Now we can run background processes
@@ -277,6 +279,10 @@ class SystemDatabase:
             with self.engine.begin() as c:
                 c.execute(cmd)
 
+        # Record we have exported status for this single-transaction workflow
+        if status["workflow_uuid"] in self._temp_txn_wf_ids:
+            self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
+
     def set_workflow_status(
         self,
         workflow_uuid: str,
@@ -467,6 +473,11 @@ class SystemDatabase:
             with self.engine.begin() as c:
                 c.execute(cmd)
 
+        if workflow_uuid in self._temp_txn_wf_ids:
+            # Clean up the single-transaction tracking sets
+            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]:
         with self.engine.begin() as c:
             row = c.execute(
@@ -546,8 +557,10 @@ class SystemDatabase:
             else:
                 with self.engine.begin() as c:
                     c.execute(sql)
-        except sa.exc.IntegrityError:
-            raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
+            raise dbapi_error
         except Exception as e:
             raise e
 
@@ -601,8 +614,11 @@ class SystemDatabase:
                         message=utils.serialize(message),
                     )
                 )
-            except sa.exc.IntegrityError:
-                raise DBOSNonExistentWorkflowError(destination_uuid)
+            except DBAPIError as dbapi_error:
+                # Foreign key violation
+                if dbapi_error.orig.pgcode == "23503":  # type: ignore
+                    raise DBOSNonExistentWorkflowError(destination_uuid)
+                raise dbapi_error
             except Exception as e:
                 raise e
             output: OperationResultInternal = {
@@ -825,8 +841,10 @@ class SystemDatabase:
                         value=utils.serialize(message),
                     )
                 )
-            except sa.exc.IntegrityError:
-                raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
+            except DBAPIError as dbapi_error:
+                if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                    raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
+                raise dbapi_error
             except Exception as e:
                 raise e
             output: OperationResultInternal = {
@@ -910,7 +928,7 @@ class SystemDatabase:
         return value
 
     def _flush_workflow_status_buffer(self) -> None:
-        """Export the workflow status buffer to the database, up to the batch size."""
+        """Export the workflow status buffer to the database, up to the batch size"""
         if len(self._workflow_status_buffer) == 0:
             return
 
@@ -918,13 +936,16 @@ class SystemDatabase:
         exported_status: Dict[str, WorkflowStatusInternal] = {}
         with self.engine.begin() as c:
             exported = 0
-            local_batch_size = min(
-                buffer_flush_batch_size, len(self._workflow_status_buffer)
-            )
-            while exported < local_batch_size:
+            status_iter = iter(list(self._workflow_status_buffer))
+            wf_id: Optional[str] = None
+            while (
+                exported < buffer_flush_batch_size
+                and (wf_id := next(status_iter, None)) is not None
+            ):
                 # Pop the first key in the buffer (FIFO)
-                wf_id = next(iter(self._workflow_status_buffer))
-                status = self._workflow_status_buffer.pop(wf_id)
+                status = self._workflow_status_buffer.pop(wf_id, None)
+                if status is None:
+                    continue
                 exported_status[wf_id] = status
                 try:
                     self.update_workflow_status(status, conn=c)
@@ -945,12 +966,18 @@ class SystemDatabase:
         exported_inputs: Dict[str, str] = {}
         with self.engine.begin() as c:
             exported = 0
-            local_batch_size = min(
-                buffer_flush_batch_size, len(self._workflow_inputs_buffer)
-            )
-            while exported < local_batch_size:
-                wf_id = next(iter(self._workflow_inputs_buffer))
-                inputs = self._workflow_inputs_buffer.pop(wf_id)
+            input_iter = iter(list(self._workflow_inputs_buffer))
+            wf_id: Optional[str] = None
+            while (
+                exported < buffer_flush_batch_size
+                and (wf_id := next(input_iter, None)) is not None
+            ):
+                if wf_id not in self._exported_temp_txn_wf_status:
+                    # Skip exporting inputs if the status has not been exported yet
+                    continue
+                inputs = self._workflow_inputs_buffer.pop(wf_id, None)
+                if inputs is None:
+                    continue
                 exported_inputs[wf_id] = inputs
                 try:
                     self.update_workflow_inputs(wf_id, inputs, conn=c)
@@ -985,6 +1012,7 @@ class SystemDatabase:
     def buffer_workflow_inputs(self, workflow_id: str, inputs: str) -> None:
         # inputs is a serialized WorkflowInputs string
         self._workflow_inputs_buffer[workflow_id] = inputs
+        self._temp_txn_wf_ids.add(workflow_id)
 
     @property
     def _is_buffers_empty(self) -> bool:

{dbos-0.5.0a3.dist-info → dbos-0.5.0a5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.5.0a3
+Version: 0.5.0a5
 Summary: A Python framework for backends that scale
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-0.5.0a3.dist-info → dbos-0.5.0a5.dist-info}/RECORD

@@ -1,18 +1,18 @@
-dbos-0.5.0a3.dist-info/METADATA,sha256=zuhmt0hFtzVr_gTyFZ2JVIbo5Z3Cu5xFPsKX6Obtsco,5420
-dbos-0.5.0a3.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
-dbos-0.5.0a3.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
-dbos-0.5.0a3.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.5.0a5.dist-info/METADATA,sha256=WRKFIucNawLOnizlrKNBp1cTcYW2BcuIcaYaQjXq1FU,5420
+dbos-0.5.0a5.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
+dbos-0.5.0a5.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
+dbos-0.5.0a5.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=X1LdP36NomDtvPfFwoMNtgXf81TO05jj7vltsp79UUw,787
 dbos/admin_sever.py,sha256=KtzH6aKyskCm4h3yulpy9jb5PIqRlYI2sjctw5mvaKY,3395
-dbos/application_database.py,sha256=cBSfyB-5KGFBtigLpFPWEAgmrKONAUnG1825vXPifBk,5411
+dbos/application_database.py,sha256=1K3kE96BgGi_QWOd2heXluyNTwFAwlUVuAR6JKKUqf0,5659
 dbos/cli.py,sha256=QnbGtZ8S963q3iyFvXNBcL4DB35r4SFMarlb5DRqN6M,7915
-dbos/context.py,sha256=_3GV1f9ztGqs-EAzTfVbfWybSfoca3BEP8iSVud-qIM,15801
-dbos/core.py,sha256=8PfbyhrIOfk2KGK8pd3Fhv3cBRT9epTkeRSVJO5UnhQ,28988
+dbos/context.py,sha256=JZMV2RtSpTK7lnyyWxeBmGPwrZSB00XZEP6R6MT9ygQ,15690
+dbos/core.py,sha256=HfKnPpIaQqIBAHzP2hD67aSIchTHp87NgD21CcujKkE,28300
 dbos/dbos-config.schema.json,sha256=azpfmoDZg7WfSy3kvIsk9iEiKB_-VZt03VEOoXJAkqE,5331
-dbos/dbos.py,sha256=HMFBO2tUii7Wp-wKeUwHPv6RPhxU422zQOumMaWI3Qg,26506
+dbos/dbos.py,sha256=HngS2BUWSbWPmloXGr-KE81BQ8dpZtlvOXM4tx4_Qhg,26246
 dbos/dbos_config.py,sha256=EkO0c0xaIM7_vAAqqnvNNEAKG5fOJbmmalqnZvaKYZA,5312
 dbos/decorators.py,sha256=lbPefsLK6Cya4cb7TrOcLglOpGT3pc6qjZdsQKlfZLg,629
-dbos/error.py,sha256=Yrs5bXmMdE8k3cux6nngQlg4DT27Cx6EmB8ddbe4-6A,3852
+dbos/error.py,sha256=nBdLC4hxGO_K9V26YbDGOo7xi1CKuN4PsE_cBv7K8Cc,3798
 dbos/fastapi.py,sha256=ZFcMizyv3pizo5zf0sSF6U4GoR3rQH8LxGipkQIGHfU,2282
 dbos/logger.py,sha256=cfybbu6F1zsgYLEPW8D8V6h033u-YedLXnGMnQQM6-4,3341
 dbos/migrations/env.py,sha256=38SIGVbmn_VV2x2u1aHLcPOoWgZ84eCymf3g_NljmbU,1626
@@ -21,14 +21,14 @@ dbos/migrations/versions/5c361fc04708_added_system_tables.py,sha256=QMgFMb0aLgC2
 dbos/migrations/versions/a3b18ad34abe_added_triggers.py,sha256=Rv0ZsZYZ_WdgGEULYsPfnp4YzaO5L198gDTgYY39AVA,2022
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 dbos/recovery.py,sha256=xfwQFWNuD8DXg5HD5_-3tG7Neo9j-x1lrqiwtn5FSh8,2015
-dbos/registrations.py,sha256=09m-w00Dh7dsXR61LjdfpLBg-7baYeaTZY-xTAD8NjM,6547
+dbos/registrations.py,sha256=gMI-u05tv5bpvyddQGtoUgCsqARx51aOY7p0JXPafQo,6539
 dbos/roles.py,sha256=9u0z4CWmXPeqIKzQWEzaOKIlzOuaagBtMiB-swqjX_U,2291
 dbos/scheduler/croniter.py,sha256=hbhgfsHBqclUS8VeLnJ9PSE9Z54z6mi4nnrr1aUXn0k,47561
 dbos/scheduler/scheduler.py,sha256=uO4_9jmWW2rLv1ODL3lc1cE_37ZaVTgnvmFx_FAlN50,1472
 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=5V3vqnEzry0Hn7ZbVS9Gs_dJKia8uX8p7mGC82Ru8rk,4303
-dbos/system_database.py,sha256=5veg8h8zw42ZBZf68myN4b1N2Z2Ri8PseJzf7IINipY,38253
+dbos/system_database.py,sha256=SK24Avj10rbbWFilVUexdPX6VvOL8zC-CoWDhNQj6QM,39698
 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=hJgp3S14cseT7zWIZsPwjqdzwTCw1aLo8kPKsTvYz0Y,2976
@@ -42,4 +42,4 @@ dbos/templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKs
 dbos/tracer.py,sha256=RPW9oxmX9tSc0Yq7O-FAhpQWBg1QT7Ni1Q06uwhtNDk,2237
 dbos/utils.py,sha256=hWj9iWDrby2cVEhb0pG-IdnrxLqP64NhkaWUXiLc8bA,402
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.5.0a3.dist-info/RECORD,,
+dbos-0.5.0a5.dist-info/RECORD,,