fix: Reject state changes from resolvers of an ended sync

A resolver still running after its sync ended could change the state client of
the next sync, e.g. one a plugin keeps on an attribute. Its key then failed every
flush of that sync, as a failed flush keeps its changes, and the next sync read
its value. set_key now raises before it records the change.

Author: jpollard-cs

cloudquery/sdk/internal/checkpoint.py

@@ -112,6 +112,20 @@ def current_scope(self) -> Optional[Scope]:
         """Returns the scope of the table resolver running in this context, or None outside table resolvers."""
         return self._scope.get()
 
+    def check_scope(self, scope: Optional[Scope]) -> None:
+        """
+        Raises if scope belongs to a sync that ended. The state client calls it before it records
+        a change, so that a resolver still running after its sync cannot change the state of the
+        next one: neither the values the next sync reads nor the keys it flushes.
+
+        Raises:
+            RuntimeError: if scope belongs to a sync that ended.
+        """
+        if scope is not None and scope.sync is not self._sync:
+            raise RuntimeError(
+                "the state of a table resolver can only be changed during its own sync"
+            )
+
     def resolver_finished(self, scope: Scope) -> None:
         """
         Marks the resolver of scope as returned, once every resource it produced is on the queue.

cloudquery/sdk/stateclient/stateclient.py

@@ -85,7 +85,10 @@ def get_key(self, key: str) -> Optional[str]:
             return self.mem.get(key)
 
     def set_key(self, key: str, value: str) -> None:
-        scope = self.checkpoints.current_scope() if self.checkpoints else None
+        scope = None
+        if self.checkpoints:
+            scope = self.checkpoints.current_scope()
+            self.checkpoints.check_scope(scope)
         with self._lock:
             self.mem[key] = value
             self.changes.setdefault(key, set()).add(scope)

tests/scheduler/checkpoints.py

@@ -395,3 +395,51 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one():
     ]
     assert cursors
     assert all(cursor.startswith("next-") for cursor in cursors)
+
+
+def test_resolvers_of_an_ended_sync_cannot_change_the_state_of_the_next_one():
+    # the plugin keeps the current sync's state client on an attribute that resolvers read
+    shared: Dict[str, StateClient] = {}
+    next_started, leftover_done = threading.Event(), threading.Event()
+    leftover_errors: List[Exception] = []
+
+    class LeftoverResolver(TableResolver):
+        def __init__(self):
+            super().__init__(Table("leftover", [Column("id", pa.int64())]))
+
+        def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+            yield {"id": 1}
+            next_started.wait(5)
+            try:
+                shared["state"].set_key("leftover_cursor", "ended")
+            except RuntimeError as e:
+                leftover_errors.append(e)
+            finally:
+                leftover_done.set()
+
+    class NextResolver(CursorResolver):
+        def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+            next_started.set()
+            leftover_done.wait(5)
+            yield from super().resolve(client, parent_resource)
+
+    scheduler = Scheduler(4)
+    try:
+        ended = scheduler.sync(TestClient(), [LeftoverResolver()])
+        while not isinstance(next(ended), SyncInsertMessage):
+            pass
+        ended.close()
+
+        shared["state"] = MemoryStateClient(checkpoints=scheduler.checkpoints)
+        scheduler.set_post_sync_hook(shared["state"].flush)
+        messages = list(
+            scheduler.sync(TestClient(), [NextResolver(shared["state"], "next")])
+        )
+    finally:
+        scheduler.shutdown()
+    assert len(leftover_errors) == 1
+    assert "leftover_cursor" not in shared["state"].mem
+    checkpoints = [
+        checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage)
+    ]
+    assert [c["cursor"] for c in checkpoints] == [f"next-{i}" for i in range(1, 30)]