diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py index 394d6a3..3e0dbf2 100644 --- a/cloudquery/sdk/internal/checkpoint.py +++ b/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" ) diff --git a/cloudquery/sdk/stateclient/stateclient.py b/cloudquery/sdk/stateclient/stateclient.py index dfa245a..af817be 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/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] diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py index 61fc182..8cf7741 100644 --- a/tests/scheduler/checkpoints.py +++ b/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"}]