Allow override the job cursors namespace

Author: leameow

src/saturn_engine/stores/jobs_store.py

@@ -143,16 +143,16 @@ def sync_jobs_states(
                     "enabled": False,
                 }
             )
+
+        namespace = job_definition_by_name.get(job_id) or job_id
         for cursor, cursor_state in job_state.cursors_states.items():
-            job_definition_name = job_definition_by_name.get(job_id)
-            if job_definition_name:
-                job_cursors.append(
-                    {
-                        "job_definition_name": job_definition_name,
-                        "cursor": cursor,
-                        "state": cursor_state,
-                    }
-                )
+            job_cursors.append(
+                {
+                    "job_definition_name": namespace,
+                    "cursor": cursor,
+                    "state": cursor_state,
+                }
+            )
 
         if job_values:
             job_values["name"] = job_id

src/saturn_engine/worker/services/job_state/service.py

@@ -97,8 +97,13 @@ async def on_items_batched(self, batch: ItemsBatch) -> None:
         cursors = [
             c for i in batch.items if (c := i.metadata["job_state"]["state_cursor"])
         ]
+        namespace = (
+            batch.job.config.get("job_state", {}).get("cursors_states_namespace")
+            or batch.job.name
+        )
+
         cursors_states: dict = await self.fetch_cursors_states(
-            batch.job.name, cursors=cursors
+            namespace, cursors=cursors
         )
         for item in batch.items:
             metadata = item.metadata.setdefault("job_state", {})
@@ -114,12 +119,20 @@ async def on_pipeline_events_emitted(self, pevents: PipelineEventsEmitted) -> No
         for event in pevents.events:
             if not isinstance(event, CursorStateUpdated):
                 continue
-            job = pevents.xmsg.queue.definition.name
+
             message = pevents.xmsg.message.message
             cursor = message.metadata.get("job_state", {}).get("state_cursor")
             if not cursor:
                 continue
-            self.set_job_cursor_state(job, cursor=cursor, cursor_state=event.state)
+
+            queue = pevents.xmsg.queue.definition
+            namespace = (
+                queue.config.get("job_state", {}).get("cursors_states_namespace")
+                or queue.name
+            )
+            self.set_job_cursor_state(
+                namespace, cursor=cursor, cursor_state=event.state
+            )
 
     def set_job_cursor(self, job_name: JobId, *, cursor: Cursor) -> None:
         self._store.set_job_cursor(job_name, cursor)

tests/worker/conftest.py

@@ -446,7 +446,9 @@ async def fetch(
 
 
 @pytest.fixture
-async def inmemory_cursors_fetcher(services_manager: ServicesManager) -> None:
+async def inmemory_cursors_fetcher(
+    services_manager: ServicesManager,
+) -> None:
     job_state_service = services_manager.services.cast_service(JobStateService)
     job_state_service._cursors_fetcher = InMemoryCursorsFetcher(
         job_state=job_state_service

tests/worker/services/job_state/test_service.py

@@ -201,17 +201,20 @@ async def test_job_state_set_message_cursor_state(
     job_state_service = services_manager.services.cast_service(JobStateService)
     fetch_cursors = mocker.spy(job_state_service, "fetch_cursors_states")
     job_state_service.set_job_cursor_state(
-        fake_queue_item.name, cursor=Cursor("1"), cursor_state={"x": 1}
+        JobId("test-job"), cursor=Cursor("1"), cursor_state={"x": 1}
     )
     job_state_service.set_job_cursor_state(
-        fake_queue_item.name, cursor=Cursor("2"), cursor_state={"x": 2}
+        JobId("test-job"), cursor=Cursor("2"), cursor_state={"x": 2}
     )
     fake_queue_item.config["job"] = {
         "buffer_flush_after": 7,
         "buffer_size": 2,
     }
     fake_queue_item.pipeline.info = PipelineInfo.from_pipeline(pipeline)
-    fake_queue_item.config["job_state"] = {"cursors_states_enabled": True}
+    fake_queue_item.config["job_state"] = {
+        "cursors_states_enabled": True,
+        "cursors_states_namespace": "test-job",
+    }
 
     @services_manager.services.s.hooks.work_queue_built.emit
     async def fake_work_builder(queue: QueueItemWithState) -> t.Any:
@@ -257,13 +260,13 @@ async def fake_work_builder(queue: QueueItemWithState) -> t.Any:
     ]
 
     assert fetch_cursors.call_args_list == [
-        call(fake_queue_item.name, cursors=[Cursor("0")]),
-        call(fake_queue_item.name, cursors=[Cursor("1"), Cursor("2")]),
-        call(fake_queue_item.name, cursors=[Cursor("3"), Cursor("a")]),
+        call("test-job", cursors=[Cursor("0")]),
+        call("test-job", cursors=[Cursor("1"), Cursor("2")]),
+        call("test-job", cursors=[Cursor("3"), Cursor("a")]),
     ]
 
     state = job_state_service._store._current_state
-    assert len(state.jobs) == 1
+    assert len(state.jobs) == 2
     job_state = state.jobs[fake_queue_item.name]
     assert job_state.cursor == "4"
     assert job_state.completion

tests/worker_manager/api/test_jobs.py

@@ -665,6 +665,7 @@ def test_sync_states(
         {"job": r.job_definition_name, "cursor": r.cursor, "state": r.state}
         for r in database_state
     ] == [
+        {"job": "orphan-test", "cursor": "a", "state": {"x": 1}},
         {"job": "test", "cursor": "a", "state": {"x": 1}},
         {"job": "test", "cursor": "b", "state": {"x": 3}},
     ]