fix: Drop state changes of ended syncs instead of rejecting them

Author: erezrokah

cloudquery/sdk/internal/checkpoint.py

@@ -112,19 +112,13 @@ 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 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:
         """
@@ -166,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,16 +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 = None
-        if self.checkpoints:
-            scope = self.checkpoints.current_scope()
-            self.checkpoints.check_scope(scope)
+        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)
@@ -98,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, {}
@@ -115,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__(
@@ -398,8 +398,7 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one():
 
 
 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] = {}
+    plugin_attributes: Dict[str, StateClient] = {}
     next_started, leftover_done = threading.Event(), threading.Event()
     leftover_errors: List[Exception] = []
 
@@ -409,18 +408,18 @@ def __init__(self):
 
         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:
+                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()
-            leftover_done.wait(5)
+            assert leftover_done.wait(5), "the leftover resolver did not finish"
             yield from super().resolve(client, parent_resource)
 
     scheduler = Scheduler(4)
@@ -430,16 +429,69 @@ def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
             pass
         ended.close()
 
-        shared["state"] = MemoryStateClient(checkpoints=scheduler.checkpoints)
-        scheduler.set_post_sync_hook(shared["state"].flush)
+        plugin_attributes["state"] = MemoryStateClient(
+            checkpoints=scheduler.checkpoints
+        )
+        scheduler.set_post_sync_hook(plugin_attributes["state"].flush)
         messages = list(
-            scheduler.sync(TestClient(), [NextResolver(shared["state"], "next")])
+            scheduler.sync(
+                TestClient(), [NextResolver(plugin_attributes["state"], "next")]
+            )
         )
     finally:
         scheduler.shutdown()
-    assert len(leftover_errors) == 1
-    assert "leftover_cursor" not in shared["state"].mem
+    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"}]