service(State): Allow to override the updated state cursor

Flared · May 28, 2024 · f1b2e4a · f1b2e4a
1 parent 7353f14
commit f1b2e4a
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git a/src/saturn_engine/core/job_state.py b/src/saturn_engine/core/job_state.py
@@ -1,11 +1,13 @@
 import dataclasses
 
+from saturn_engine.core import Cursor
 from saturn_engine.core.pipeline import PipelineEvent
 
 
 @dataclasses.dataclass
 class CursorStateUpdated(PipelineEvent):
     state: dict
+    cursor: Cursor | None = None
 
 
 class CursorState:

diff --git a/src/saturn_engine/worker/services/job_state/service.py b/src/saturn_engine/worker/services/job_state/service.py
@@ -198,7 +198,9 @@ async def on_pipeline_events_emitted(self, pevents: PipelineEventsEmitted) -> No
                 continue
 
             message = pevents.xmsg.message.message
-            cursor = message.metadata.get("job_state", {}).get("state_cursor")
+            cursor = event.cursor or message.metadata.get("job_state", {}).get(
+                "state_cursor"
+            )
             if not cursor:
                 continue
 

diff --git a/tests/worker/services/job_state/test_service.py b/tests/worker/services/job_state/test_service.py
@@ -253,12 +253,22 @@ async def fake_work_builder(queue: QueueItemWithState) -> t.Any:
         async with xmsg._context:
             xmsgs.append(xmsg)
             results.append(xmsg.message.execute())
+
+            # Emit a new state.
             await services_manager.services.s.hooks.pipeline_events_emitted.emit(
                 PipelineEventsEmitted(
                     xmsg=xmsg, events=[CursorStateUpdated(state={"x": len(xmsgs) * 10})]
                 )
             )
 
+            # Emit a new state with a custom cursor.
+            await services_manager.services.s.hooks.pipeline_events_emitted.emit(
+                PipelineEventsEmitted(
+                    xmsg=xmsg,
+                    events=[CursorStateUpdated(cursor=Cursor("42"), state={"x": 42})],
+                )
+            )
+
     # Assert everything loaded and ran in order.
     msgs = [i.message.message for i in xmsgs]
     assert [i.id for i in msgs] == ["0", "1", "2", "3", "4"]
@@ -287,7 +297,14 @@ async def fake_work_builder(queue: QueueItemWithState) -> t.Any:
     # Rerun the inventory, expect new states to be loaded
     inventory = FakeInventory(
         options=FakeInventory.Options(
-            data=[5, 4, 3, 2, (1, {"job_state": {"state_cursor": "a"}})]
+            data=[
+                5,
+                4,
+                3,
+                2,
+                (1, {"job_state": {"state_cursor": "a"}}),
+                (42, {"job_state": {"state_cursor": "42"}}),
+            ]
         )
     )
     job = Job(
@@ -310,6 +327,7 @@ async def fake_work_builder(queue: QueueItemWithState) -> t.Any:
         {"x": 30},
         {"x": 40},
         {"x": 50},
+        {"x": 42},
     ]