Skip to content

Commit 0a01d27

Browse files
committed
Merge remote-tracking branch 'origin/main' into fix/checkpoint-review-followups
# Conflicts: # cloudquery/sdk/internal/checkpoint.py # cloudquery/sdk/stateclient/stateclient.py # tests/scheduler/checkpoints.py
2 parents d220d2f + 0a4b408 commit 0a01d27

3 files changed

Lines changed: 134 additions & 3 deletions

File tree

‎cloudquery/sdk/internal/checkpoint.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,14 @@ def resource(self, scope: Scope) -> Iterator[None]:
147147
scope.after_resource = []
148148
self._send(ready)
149149

150+
def superseded(self, scope: Optional[Scope]) -> bool:
151+
"""Returns True if scope belongs to a sync that ended while another sync runs now."""
152+
return (
153+
scope is not None
154+
and self._sync is not None
155+
and scope.sync is not self._sync
156+
)
157+
150158
def resolver_finished(self, scope: Scope) -> None:
151159
"""
152160
Marks the resolver of scope as returned, once every resource it produced is on the queue.
@@ -187,7 +195,7 @@ def schedule(
187195
"e.g. from a table resolver or the post-sync hook"
188196
)
189197
for scopes, _, _ in checkpoints:
190-
if any(s is not None and s.sync is not self._sync for s in scopes):
198+
if any(self.superseded(s) for s in scopes):
191199
raise SyncNotRunningError(
192200
"state checkpoints of a table resolver can only be flushed during its own sync"
193201
)

‎cloudquery/sdk/stateclient/stateclient.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,13 +79,17 @@ def __init__(
7979

8080
self.migrate_state_table()
8181
self.read_all_state()
82+
self._flushed = dict(self.mem)
8283

8384
def get_key(self, key: str) -> Optional[str]:
8485
with self._lock:
86+
self._forget_changes_of_ended_syncs(key)
8587
return self.mem.get(key)
8688

8789
def set_key(self, key: str, value: str) -> None:
8890
scope = self.checkpoints.current_scope() if self.checkpoints else None
91+
if self.checkpoints and self.checkpoints.superseded(scope):
92+
return
8993
with self._lock:
9094
if self.mem.get(key) == value:
9195
# a pending change of the key summarizes this caller's resources too
@@ -100,13 +104,17 @@ def flush(self):
100104
Persists the keys changed since the last flush: writes them to the backend, or, with
101105
checkpoints, sends them on the sync stream after the resources they cover.
102106
If that fails, the keys stay changed, so a later flush retries them.
107+
Changes made by table resolvers of an ended sync are dropped once another sync runs,
108+
as the resources they cover may never have been committed.
103109
104110
Raises:
105111
SyncNotRunningError: with checkpoints, if the sync is not running, e.g. once
106112
Scheduler.sync returned. Flush in the scheduler's post-sync hook instead.
107113
"""
108114
with self._flush_lock:
109115
with self._lock:
116+
for key in list(self.changes):
117+
self._forget_changes_of_ended_syncs(key)
110118
if not self.changes:
111119
return
112120
changes, self.changes = self.changes, {}
@@ -121,6 +129,21 @@ def flush(self):
121129
for key, scopes in changes.items():
122130
self.changes.setdefault(key, set()).update(scopes)
123131
raise
132+
with self._lock:
133+
self._flushed.update(values)
134+
135+
def _forget_changes_of_ended_syncs(self, key: str) -> None:
136+
if self.checkpoints is None or key not in self.changes:
137+
return
138+
scopes = {s for s in self.changes[key] if not self.checkpoints.superseded(s)}
139+
if scopes:
140+
self.changes[key] = scopes
141+
return
142+
del self.changes[key]
143+
if key in self._flushed:
144+
self.mem[key] = self._flushed[key]
145+
else:
146+
self.mem.pop(key, None)
124147

125148
def _send_checkpoints(
126149
self, changes: Dict[str, Set[Optional[Scope]]], values: Dict[str, str]

‎tests/scheduler/checkpoints.py‎

Lines changed: 102 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@
3131
class MemoryStateClient(StateClient):
3232
"""A state client whose backend is a dict."""
3333

34-
def __init__(self, checkpoints=None, fail_writes=False):
35-
self.backend: Dict[str, str] = {}
34+
def __init__(self, checkpoints=None, fail_writes=False, backend=None):
35+
self.backend: Dict[str, str] = dict(backend or {})
3636
self.writes: List[List[Tuple[str, str]]] = []
3737
self.fail_writes = fail_writes
3838
super().__init__(
@@ -423,6 +423,106 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one():
423423
assert all(cursor.startswith("next-") for cursor in cursors)
424424

425425

426+
def test_resolvers_of_an_ended_sync_cannot_change_the_state_of_the_next_one():
427+
plugin_attributes: Dict[str, StateClient] = {}
428+
next_started, leftover_done = threading.Event(), threading.Event()
429+
leftover_errors: List[Exception] = []
430+
431+
class LeftoverResolver(TableResolver):
432+
def __init__(self):
433+
super().__init__(Table("leftover", [Column("id", pa.int64())]))
434+
435+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
436+
yield {"id": 1}
437+
try:
438+
assert next_started.wait(5), "the next sync did not start"
439+
plugin_attributes["state"].set_key("leftover_cursor", "ended")
440+
except Exception as e:
441+
leftover_errors.append(e)
442+
finally:
443+
leftover_done.set()
444+
445+
class NextResolver(CursorResolver):
446+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
447+
next_started.set()
448+
assert leftover_done.wait(5), "the leftover resolver did not finish"
449+
yield from super().resolve(client, parent_resource)
450+
451+
scheduler = Scheduler(4)
452+
try:
453+
ended = scheduler.sync(TestClient(), [LeftoverResolver()])
454+
while not isinstance(next(ended), SyncInsertMessage):
455+
pass
456+
ended.close()
457+
458+
plugin_attributes["state"] = MemoryStateClient(
459+
checkpoints=scheduler.checkpoints
460+
)
461+
scheduler.set_post_sync_hook(plugin_attributes["state"].flush)
462+
messages = list(
463+
scheduler.sync(
464+
TestClient(), [NextResolver(plugin_attributes["state"], "next")]
465+
)
466+
)
467+
finally:
468+
scheduler.shutdown()
469+
assert leftover_errors == []
470+
assert plugin_attributes["state"].get_key("leftover_cursor") is None
471+
checkpoints = [
472+
checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage)
473+
]
474+
assert [c["cursor"] for c in checkpoints] == [f"next-{i}" for i in range(1, 30)]
475+
476+
477+
class SetKeyThenWait(TableResolver):
478+
def __init__(self, state: StateClient, release: threading.Event):
479+
super().__init__(Table("set_then_wait", [Column("id", pa.int64())]))
480+
self._state = state
481+
self._release = release
482+
483+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
484+
self._state.set_key("cursor", "ended")
485+
yield {"id": 1}
486+
self._release.wait(5)
487+
488+
489+
def test_a_state_client_reused_across_syncs_drops_unflushed_changes_of_an_ended_sync():
490+
release = threading.Event()
491+
scheduler = Scheduler(4)
492+
try:
493+
state = MemoryStateClient(
494+
checkpoints=scheduler.checkpoints, backend={"cursor": "flushed"}
495+
)
496+
ended = scheduler.sync(TestClient(), [SetKeyThenWait(state, release)])
497+
while not isinstance(next(ended), SyncInsertMessage):
498+
pass
499+
ended.close()
500+
release.set()
501+
502+
scheduler.set_post_sync_hook(state.flush)
503+
cursors_seen: List[Optional[str]] = []
504+
505+
class ReadCursor(TableResolver):
506+
def __init__(self):
507+
super().__init__(Table("read_cursor", [Column("id", pa.int64())]))
508+
509+
def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
510+
cursors_seen.append(state.get_key("cursor"))
511+
yield {"id": 1}
512+
state.set_key("other", "next")
513+
514+
messages = list(scheduler.sync(TestClient(), [ReadCursor()]))
515+
list(scheduler.sync(TestClient(), [ReadCursor()]))
516+
finally:
517+
release.set()
518+
scheduler.shutdown()
519+
assert cursors_seen == ["flushed", "flushed"]
520+
checkpoints = [
521+
checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage)
522+
]
523+
assert checkpoints == [{"other": "next"}]
524+
525+
426526
class HookCursorResolver(TableResolver):
427527
"""Checkpoints its cursor from a resource hook, before the resource is on the queue."""
428528

0 commit comments

Comments
 (0)