fix: Raise SyncNotRunningError for checkpoints outside their sync, and skip unchanged state values

Author: disq

cloudquery/sdk/internal/checkpoint.py

@@ -22,6 +22,14 @@
 from cloudquery.sdk.message import SyncCheckpointMessage
 
 
+class SyncNotRunningError(RuntimeError):
+    """
+    Raised when state checkpoints are flushed while their sync is not running, so they cannot be
+    sent: before or after Scheduler.sync, or from a resolver of an ended sync. Flush state that
+    summarizes the whole sync in the scheduler's post-sync hook instead.
+    """
+
+
 class Scope:
     """
     A table resolver call and the relations its resources spawn.
@@ -170,17 +178,17 @@ def schedule(
         wait for nothing. Schedules all of them or, if it raises, none of them.
 
         Raises:
-            RuntimeError: if no sync is running, as the checkpoints could not be sent.
+            SyncNotRunningError: if the sync of the checkpoints is not running.
         """
         with self._lock:
             if self._res is None:
-                raise RuntimeError(
+                raise SyncNotRunningError(
                     "state checkpoints can only be flushed while the scheduler syncs, "
                     "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):
-                    raise RuntimeError(
+                    raise SyncNotRunningError(
                         "state checkpoints of a table resolver can only be flushed during its own sync"
                     )
             for scopes, keys, msg in checkpoints:

cloudquery/sdk/stateclient/__init__.py

@@ -1 +1,2 @@
+from cloudquery.sdk.internal.checkpoint import SyncNotRunningError
 from .keys import state_key

cloudquery/sdk/stateclient/stateclient.py

@@ -87,6 +87,11 @@ def get_key(self, key: str) -> Optional[str]:
     def set_key(self, key: str, value: str) -> None:
         scope = self.checkpoints.current_scope() if self.checkpoints else None
         with self._lock:
+            if self.mem.get(key) == value:
+                # a pending change of the key summarizes this caller's resources too
+                if key in self.changes:
+                    self.changes[key].add(scope)
+                return
             self.mem[key] = value
             self.changes.setdefault(key, set()).add(scope)
 
@@ -95,6 +100,10 @@ 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.
+
+        Raises:
+            SyncNotRunningError: with checkpoints, if the sync is not running, e.g. once
+                Scheduler.sync returned. Flush in the scheduler's post-sync hook instead.
         """
         with self._flush_lock:
             with self._lock:

tests/scheduler/checkpoints.py

@@ -18,6 +18,7 @@
 from cloudquery.sdk.schema import Column
 from cloudquery.sdk.schema.arrow import METADATA_TABLE_NAME
 from cloudquery.sdk.schema.table import Table
+from cloudquery.sdk.stateclient import SyncNotRunningError
 from cloudquery.sdk.stateclient.stateclient import (
     ConnectedStateClient,
     StateClient,
@@ -223,6 +224,31 @@ def test_flush_writes_to_the_backend_without_checkpoints():
     assert state.writes[-1] == [("a", "3")]
 
 
+def test_setting_an_unchanged_value_is_not_a_change():
+    state = MemoryStateClient(checkpoints=None)
+    state.backend["a"] = "1"
+    state.read_all_state()
+    state.set_key("a", "1")
+    state.flush()
+    assert state.writes == []
+
+    state.set_key("a", "2")
+    state.set_key("a", "2")
+    state.flush()
+    assert state.writes == [[("a", "2")]]
+
+
+def test_setting_an_unchanged_value_adds_the_caller_to_the_pending_change():
+    scheduler = Scheduler(1)
+    state = MemoryStateClient(checkpoints=scheduler.checkpoints)
+    state.set_key("a", "1")
+    scope = scheduler.checkpoints.new_scope(None)
+    with scheduler.checkpoints.entered(scope):
+        state.set_key("a", "1")
+    scheduler.shutdown()
+    assert state.changes["a"] == {None, scope}
+
+
 def test_failed_flush_keeps_the_changes():
     state = MemoryStateClient(checkpoints=None, fail_writes=True)
     state.set_key("a", "1")
@@ -238,7 +264,7 @@ def test_checkpoint_flush_outside_a_sync_fails_and_keeps_the_changes():
     scheduler = Scheduler(1)
     state = MemoryStateClient(checkpoints=scheduler.checkpoints)
     state.set_key("a", "1")
-    with pytest.raises(RuntimeError):
+    with pytest.raises(SyncNotRunningError):
         state.flush()
     assert "a" in state.changes
     assert state.writes == []