Set Event From Step (#483)

Author: kraftp

dbos/_core.py

@@ -1278,21 +1278,24 @@ def recv(dbos: "DBOS", topic: Optional[str] = None, timeout_seconds: float = 60)
 def set_event(dbos: "DBOS", key: str, value: Any) -> None:
     cur_ctx = get_local_dbos_context()
     if cur_ctx is not None:
-        # Must call it within a workflow
-        assert (
-            cur_ctx.is_workflow()
-        ), "set_event() must be called from within a workflow"
-        attributes: TracedAttributes = {
-            "name": "set_event",
-        }
-        with EnterDBOSStep(attributes):
-            ctx = assert_current_dbos_context()
-            dbos._sys_db.set_event(
-                ctx.workflow_id, ctx.curr_step_function_id, key, value
+        if cur_ctx.is_workflow():
+            # If called from a workflow function, run as a step
+            attributes: TracedAttributes = {
+                "name": "set_event",
+            }
+            with EnterDBOSStep(attributes):
+                ctx = assert_current_dbos_context()
+                dbos._sys_db.set_event_from_workflow(
+                    ctx.workflow_id, ctx.curr_step_function_id, key, value
+                )
+        elif cur_ctx.is_step():
+            dbos._sys_db.set_event_from_step(cur_ctx.workflow_id, key, value)
+        else:
+            raise DBOSException(
+                "set_event() must be called from within a workflow or step"
             )
     else:
-        # Cannot call it from outside of a workflow
-        raise DBOSException("set_event() must be called from within a workflow")
+        raise DBOSException("set_event() must be called from within a workflow or step")
 
 
 def get_event(

dbos/_sys_db.py

@@ -1524,7 +1524,7 @@ def sleep(
         return duration
 
     @db_retry()
-    def set_event(
+    def set_event_from_workflow(
         self,
         workflow_uuid: str,
         function_id: int,
@@ -1566,6 +1566,26 @@ def set_event(
             }
             self._record_operation_result_txn(output, conn=c)
 
+    def set_event_from_step(
+        self,
+        workflow_uuid: str,
+        key: str,
+        message: Any,
+    ) -> None:
+        with self.engine.begin() as c:
+            c.execute(
+                self.dialect.insert(SystemSchema.workflow_events)
+                .values(
+                    workflow_uuid=workflow_uuid,
+                    key=key,
+                    value=_serialization.serialize(message),
+                )
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid", "key"],
+                    set_={"value": _serialization.serialize(message)},
+                )
+            )
+
     def get_all_events(self, workflow_id: str) -> Dict[str, Any]:
         """
         Get all events currently present for a workflow ID.

tests/test_dbos.py

@@ -965,41 +965,57 @@ def test_send_recv_workflow(topic: str) -> str:
 def test_set_get_events(dbos: DBOS) -> None:
     @DBOS.workflow()
     def test_setevent_workflow() -> None:
-        dbos.set_event("key1", "value1")
-        dbos.set_event("key2", "value2")
-        dbos.set_event("key3", None)
+        DBOS.set_event("key1", "value1")
+        DBOS.set_event("key2", "value2")
+        DBOS.set_event("key3", None)
+        set_event_step()
+
+    @DBOS.step()
+    def set_event_step() -> None:
+        DBOS.set_event("key4", "value4")
 
     @DBOS.workflow()
     def test_getevent_workflow(
-        target_uuid: str, key: str, timeout_seconds: float = 10
+        target_uuid: str, key: str, timeout: float = 0.0
     ) -> Optional[str]:
-        msg = dbos.get_event(target_uuid, key, timeout_seconds)
+        msg = dbos.get_event(target_uuid, key, timeout)
         return str(msg) if msg is not None else None
 
-    wfuuid = str(uuid.uuid4())
-    with SetWorkflowID(wfuuid):
+    wfid = str(uuid.uuid4())
+    with SetWorkflowID(wfid):
         test_setevent_workflow()
-    with SetWorkflowID(wfuuid):
+    with SetWorkflowID(wfid):
         test_setevent_workflow()
 
-    value1 = test_getevent_workflow(wfuuid, "key1")
+    value1 = test_getevent_workflow(wfid, "key1")
     assert value1 == "value1"
 
-    value2 = test_getevent_workflow(wfuuid, "key2")
+    value2 = test_getevent_workflow(wfid, "key2")
     assert value2 == "value2"
 
     # Run getEvent outside of a workflow
-    value1 = dbos.get_event(wfuuid, "key1")
+    value1 = DBOS.get_event(wfid, "key1", 0)
     assert value1 == "value1"
 
-    value2 = dbos.get_event(wfuuid, "key2")
+    value2 = DBOS.get_event(wfid, "key2", 0)
     assert value2 == "value2"
 
     begin_time = time.time()
-    value3 = test_getevent_workflow(wfuuid, "key3")
+    value3 = test_getevent_workflow(wfid, "key3")
     assert value3 is None
-    duration = time.time() - begin_time
-    assert duration < 1  # None is from the event not from the timeout
+
+    value4 = DBOS.get_event(wfid, "key4", 0)
+    assert value4 == "value4"
+
+    steps = DBOS.list_workflow_steps(wfid)
+    assert len(steps) == 4
+    assert (
+        steps[0]["function_name"]
+        == steps[1]["function_name"]
+        == steps[2]["function_name"]
+        == "DBOS.setEvent"
+    )
+    assert steps[3]["function_name"] == set_event_step.__qualname__
 
     # Test OAOO
     timeout_uuid = str(uuid.uuid4())