Set Event Multiple Times (#108)

You can now set the same event multiple times from a workflow with
`set_event`. `get_event` will retrieve the latest value of the event.

Author: kraftp

dbos/error.py

@@ -32,7 +32,6 @@ class DBOSErrorCode(Enum):
     InitializationError = 3
     WorkflowFunctionNotFound = 4
     NonExistentWorkflowError = 5
-    DuplicateWorkflowEventError = 6
     MaxStepRetriesExceeded = 7
     NotAuthorized = 8
 
@@ -87,16 +86,6 @@ def __init__(self, destination_id: str):
         )
 
 
-class DBOSDuplicateWorkflowEventError(DBOSException):
-    """Exception raised when a workflow attempts to set an event value more than once per key."""
-
-    def __init__(self, workflow_id: str, key: str):
-        super().__init__(
-            f"Workflow {workflow_id} has already emitted an event with key {key}",
-            dbos_error_code=DBOSErrorCode.DuplicateWorkflowEventError.value,
-        )
-
-
 class DBOSNotAuthorizedError(DBOSException):
     """Exception raised by DBOS role-based security when the user is not authorized to access a function."""
 

dbos/system_database.py

@@ -1,6 +1,5 @@
 import datetime
 import os
-import select
 import threading
 import time
 from enum import Enum
@@ -14,11 +13,7 @@
 from sqlalchemy.exc import DBAPIError
 
 import dbos.utils as utils
-from dbos.error import (
-    DBOSDuplicateWorkflowEventError,
-    DBOSNonExistentWorkflowError,
-    DBOSWorkflowConflictIDError,
-)
+from dbos.error import DBOSNonExistentWorkflowError, DBOSWorkflowConflictIDError
 
 from .dbos_config import ConfigFile
 from .logger import dbos_logger
@@ -832,18 +827,18 @@ def set_event(
             if recorded_output is not None:
                 return  # Already sent before
 
-            try:
-                c.execute(
-                    pg.insert(SystemSchema.workflow_events).values(
-                        workflow_uuid=workflow_uuid,
-                        key=key,
-                        value=utils.serialize(message),
-                    )
+            c.execute(
+                pg.insert(SystemSchema.workflow_events)
+                .values(
+                    workflow_uuid=workflow_uuid,
+                    key=key,
+                    value=utils.serialize(message),
                 )
-            except DBAPIError as dbapi_error:
-                if dbapi_error.orig.sqlstate == "23505":  # type: ignore
-                    raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
-                raise
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid", "key"],
+                    set_={"value": utils.serialize(message)},
+                )
+            )
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,

tests/test_dbos.py

@@ -1,7 +1,8 @@
 import datetime
+import threading
 import time
 import uuid
-from typing import Any, Optional
+from typing import Optional
 
 import pytest
 import sqlalchemy as sa
@@ -13,7 +14,6 @@
 from dbos.context import assert_current_dbos_context, get_local_dbos_context
 from dbos.error import DBOSMaxStepRetriesExceeded
 from dbos.system_database import GetWorkflowsInput
-from tests.conftest import default_config
 
 
 def test_simple_workflow(dbos: DBOS) -> None:
@@ -882,3 +882,23 @@ def test_getevent_workflow(
     with pytest.raises(Exception) as exc_info:
         dbos.set_event("key1", "value1")
     assert "set_event() must be called from within a workflow" in str(exc_info.value)
+
+
+def test_multi_set_event(dbos: DBOS) -> None:
+    event = threading.Event()
+
+    wfid = str(uuid.uuid4())
+
+    @DBOS.workflow()
+    def test_setevent_workflow() -> None:
+        assert DBOS.workflow_id == wfid
+        DBOS.set_event("key", "value1")
+        event.wait()
+        DBOS.set_event("key", "value2")
+
+    with SetWorkflowID(wfid):
+        handle = DBOS.start_workflow(test_setevent_workflow)
+    assert DBOS.get_event(wfid, "key") == "value1"
+    event.set()
+    assert handle.get_result() == None
+    assert DBOS.get_event(wfid, "key") == "value2"

tests/test_fastapi_roles.py

@@ -18,7 +18,7 @@
 
 # Private API because this is a unit test
 from dbos.context import assert_current_dbos_context
-from dbos.error import DBOSDuplicateWorkflowEventError, DBOSNotAuthorizedError
+from dbos.error import DBOSInitializationError, DBOSNotAuthorizedError
 from dbos.system_database import GetWorkflowsInput
 from dbos.tracer import dbos_tracer
 
@@ -54,7 +54,7 @@ def test_dbos_error() -> None:
 
     @app.get("/dbosinternalerror")
     def test_dbos_error_internal() -> None:
-        raise DBOSDuplicateWorkflowEventError("nosuchwf", "test")
+        raise DBOSInitializationError("oh no")
 
     @app.get("/open/{var1}")
     @DBOS.required_roles([])