fix: Keep checkpoint order for unchanged state values

Author: disq

cloudquery/sdk/stateclient/stateclient.py

@@ -91,10 +91,9 @@ def set_key(self, key: str, value: str) -> None:
         if self.checkpoints and self.checkpoints.superseded(scope):
             return
         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)
+            # with checkpoints, the value also covers the caller's resources, so later checkpoints
+            # of the key must still follow them even if it is unchanged
+            if self.checkpoints is None and self.mem.get(key) == value:
                 return
             self.mem[key] = value
             self.changes.setdefault(key, set()).add(scope)

tests/scheduler/checkpoints.py

@@ -1,4 +1,5 @@
 import contextvars
+import queue
 import random
 import threading
 import time
@@ -249,6 +250,37 @@ def test_setting_an_unchanged_value_adds_the_caller_to_the_pending_change():
     assert state.changes["a"] == {None, scope}
 
 
+def test_unchanged_values_order_later_checkpoints_after_the_callers_resources():
+    scheduler = Scheduler(1)
+    checkpoints = scheduler.checkpoints
+    res = queue.Queue()
+    checkpoints.start(res)
+    state = MemoryStateClient(checkpoints=checkpoints)
+    first, unfinished, last = (checkpoints.new_scope(None) for _ in range(3))
+    child = checkpoints.new_scope(unfinished)
+
+    def flush_in(scope, value):
+        with checkpoints.entered(scope):
+            state.set_key("cursor", value)
+            state.flush()
+
+    def sent():
+        return [checkpoint_rows(res.get_nowait())["cursor"] for _ in range(res.qsize())]
+
+    try:
+        flush_in(first, "1")
+        assert sent() == ["1"]
+        # the same value, but it also covers the resources of a resolver whose relations still run
+        flush_in(unfinished, "1")
+        flush_in(last, "2")
+        assert sent() == [], "a later value must not overtake the relations it follows"
+        checkpoints.resolver_finished(child)
+        assert sent() == ["1", "2"]
+    finally:
+        checkpoints.stop()
+        scheduler.shutdown()
+
+
 def test_failed_flush_keeps_the_changes():
     state = MemoryStateClient(checkpoints=None, fail_writes=True)
     state.set_key("a", "1")