fix: Reject state changes from resolvers of an ended sync

cloudquery/sdk/internal/checkpoint.py

@@ -112,6 +112,14 @@ 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 superseded(self, scope: Optional[Scope]) -> bool:
+        """Returns True if scope belongs to a sync that ended while another sync runs now."""
+        return (
+            scope is not None
+            and self._sync is not None
+            and scope.sync is not self._sync
+        )
+
     def resolver_finished(self, scope: Scope) -> None:
         """
         Marks the resolver of scope as returned, once every resource it produced is on the queue.
@@ -152,7 +160,7 @@ def schedule(
                     "e.g. from a table resolver or the post-sync hook"
                 )
             for scopes, _, _ in checkpoints:
-                if any(s is not None and s.sync is not self._sync for s in scopes):
+                if any(self.superseded(s) for s in scopes):
                     raise RuntimeError(
                         "state checkpoints of a table resolver can only be flushed during its own sync"
                     )

cloudquery/sdk/stateclient/stateclient.py

@@ -79,13 +79,17 @@ def __init__(
 
         self.migrate_state_table()
         self.read_all_state()
+        self._flushed = dict(self.mem)
 
     def get_key(self, key: str) -> Optional[str]:
         with self._lock:
+            self._forget_changes_of_ended_syncs(key)
             return self.mem.get(key)
 
     def set_key(self, key: str, value: str) -> None:
         scope = self.checkpoints.current_scope() if self.checkpoints else None
+        if self.checkpoints and self.checkpoints.superseded(scope):
+            return
         with self._lock:
             self.mem[key] = value
             self.changes.setdefault(key, set()).add(scope)
@@ -95,9 +99,13 @@ def flush(self):
         Persists the keys changed since the last flush: writes them to the backend, or, with
         checkpoints, sends them on the sync stream after the resources they cover.
         If that fails, the keys stay changed, so a later flush retries them.
+        Changes made by table resolvers of an ended sync are dropped once another sync runs,
+        as the resources they cover may never have been committed.
         """
         with self._flush_lock:
             with self._lock:
+                for key in list(self.changes):
+                    self._forget_changes_of_ended_syncs(key)
                 if not self.changes:
                     return
                 changes, self.changes = self.changes, {}
@@ -112,6 +120,21 @@ def flush(self):
                     for key, scopes in changes.items():
                         self.changes.setdefault(key, set()).update(scopes)
                 raise
+            with self._lock:
+                self._flushed.update(values)
+
+    def _forget_changes_of_ended_syncs(self, key: str) -> None:
+        if self.checkpoints is None or key not in self.changes:
+            return
+        scopes = {s for s in self.changes[key] if not self.checkpoints.superseded(s)}
+        if scopes:
+            self.changes[key] = scopes
+            return
+        del self.changes[key]
+        if key in self._flushed:
+            self.mem[key] = self._flushed[key]
+        else:
+            self.mem.pop(key, None)
 
     def _send_checkpoints(
         self, changes: Dict[str, Set[Optional[Scope]]], values: Dict[str, str]

tests/scheduler/checkpoints.py

@@ -30,8 +30,8 @@
 class MemoryStateClient(StateClient):
     """A state client whose backend is a dict."""
 
-    def __init__(self, checkpoints=None, fail_writes=False):
-        self.backend: Dict[str, str] = {}
+    def __init__(self, checkpoints=None, fail_writes=False, backend=None):
+        self.backend: Dict[str, str] = dict(backend or {})
         self.writes: List[List[Tuple[str, str]]] = []
         self.fail_writes = fail_writes
         super().__init__(
@@ -395,3 +395,103 @@ 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():
+    plugin_attributes: 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}
+            try:
+                assert next_started.wait(5), "the next sync did not start"
+                plugin_attributes["state"].set_key("leftover_cursor", "ended")
+            except Exception 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()
+            assert leftover_done.wait(5), "the leftover resolver did not finish"
+            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()
+
+        plugin_attributes["state"] = MemoryStateClient(
+            checkpoints=scheduler.checkpoints
+        )
+        scheduler.set_post_sync_hook(plugin_attributes["state"].flush)
+        messages = list(
+            scheduler.sync(
+                TestClient(), [NextResolver(plugin_attributes["state"], "next")]
+            )
+        )
+    finally:
+        scheduler.shutdown()
+    assert leftover_errors == []
+    assert plugin_attributes["state"].get_key("leftover_cursor") is None
+    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)]
+
+
+class SetKeyThenWait(TableResolver):
+    def __init__(self, state: StateClient, release: threading.Event):
+        super().__init__(Table("set_then_wait", [Column("id", pa.int64())]))
+        self._state = state
+        self._release = release
+
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+        self._state.set_key("cursor", "ended")
+        yield {"id": 1}
+        self._release.wait(5)
+
+
+def test_a_state_client_reused_across_syncs_drops_unflushed_changes_of_an_ended_sync():
+    release = threading.Event()
+    scheduler = Scheduler(4)
+    try:
+        state = MemoryStateClient(
+            checkpoints=scheduler.checkpoints, backend={"cursor": "flushed"}
+        )
+        ended = scheduler.sync(TestClient(), [SetKeyThenWait(state, release)])
+        while not isinstance(next(ended), SyncInsertMessage):
+            pass
+        ended.close()
+        release.set()
+
+        scheduler.set_post_sync_hook(state.flush)
+        cursors_seen: List[Optional[str]] = []
+
+        class ReadCursor(TableResolver):
+            def __init__(self):
+                super().__init__(Table("read_cursor", [Column("id", pa.int64())]))
+
+            def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+                cursors_seen.append(state.get_key("cursor"))
+                yield {"id": 1}
+                state.set_key("other", "next")
+
+        messages = list(scheduler.sync(TestClient(), [ReadCursor()]))
+        list(scheduler.sync(TestClient(), [ReadCursor()]))
+    finally:
+        release.set()
+        scheduler.shutdown()
+    assert cursors_seen == ["flushed", "flushed"]
+    checkpoints = [
+        checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage)
+    ]
+    assert checkpoints == [{"other": "next"}]