From 1b4fde9a07ced60ae88b0f4e707ecee83c277f67 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 14:17:34 +0300 Subject: [PATCH 1/2] fix: Address checkpoint review findings - Wait for _sync to finish submitting resolvers before ending the sync - Yield every message left after the post-sync hook - Track the resolver scope in a ContextVar so copied contexts keep it - Schedule a flush's checkpoints all or none - Accept non-string client ids in state_key - Drop unused _Pending.sent and the redundant root scope in _sync --- cloudquery/sdk/internal/checkpoint.py | 82 +++++++++++++---------- cloudquery/sdk/scheduler/scheduler.py | 59 +++++++--------- cloudquery/sdk/stateclient/keys.py | 4 +- cloudquery/sdk/stateclient/stateclient.py | 15 ++++- tests/scheduler/checkpoints.py | 52 ++++++++++++++ tests/stateclient/keys.py | 4 ++ 6 files changed, 144 insertions(+), 72 deletions(-) diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py index 1b04b0e..df038f0 100644 --- a/cloudquery/sdk/internal/checkpoint.py +++ b/cloudquery/sdk/internal/checkpoint.py @@ -13,10 +13,11 @@ from __future__ import annotations +import contextvars import queue import threading from contextlib import contextmanager -from typing import Dict, Iterable, Iterator, List, Optional, Set +from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple from cloudquery.sdk.message import SyncCheckpointMessage @@ -49,7 +50,6 @@ def __init__(self, msg: SyncCheckpointMessage, keys: List[str]): self.keys = keys self.blockers = 0 self.dependents: List[_Pending] = [] - self.sent = False class Checkpoints: @@ -62,7 +62,9 @@ class Checkpoints: def __init__(self): self._lock = threading.Lock() - self._local = threading.local() + self._scope: contextvars.ContextVar[Optional[Scope]] = contextvars.ContextVar( + "cloudquery_checkpoint_scope", default=None + ) self._res: Optional[queue.Queue] = None # newest unsent checkpoint of each key self._unsent: Dict[str, _Pending] = {} @@ -89,17 +91,20 @@ def new_scope(self, parent: Optional[Scope]) -> Scope: @contextmanager def entered(self, scope: Scope) -> Iterator[None]: - """Makes scope the current one on this thread, for the state changes its resolver makes.""" - previous = getattr(self._local, "scope", None) - self._local.scope = scope + """ + Makes scope the current one in this context, for the state changes its resolver makes. + To change state from other threads, run them in a copy of the context, e.g. with + ``contextvars.copy_context().run``; asyncio tasks copy it already. + """ + token = self._scope.set(scope) try: yield finally: - self._local.scope = previous + self._scope.reset(token) def current_scope(self) -> Optional[Scope]: - """Returns the scope of the table resolver running on this thread, or None outside table resolvers.""" - return getattr(self._local, "scope", None) + """Returns the scope of the table resolver running in this context, or None outside table resolvers.""" + return self._scope.get() def resolver_finished(self, scope: Scope) -> None: """ @@ -121,17 +126,18 @@ def resolver_finished(self, scope: Scope) -> None: def schedule( self, - scopes: Iterable[Optional[Scope]], - keys: List[str], - msg: SyncCheckpointMessage, + checkpoints: List[ + Tuple[Iterable[Optional[Scope]], List[str], SyncCheckpointMessage] + ], ) -> None: """ - Sends msg once every resource the scopes produced so far is on the queue, including - the relations spawned from them, and after every earlier unsent checkpoint of keys. - A None scope stands for changes made outside table resolvers, which wait for nothing. + Sends each (scopes, keys, msg) once every resource the scopes produced so far is on the + queue, including the relations spawned from them, and after every earlier unsent + checkpoint of keys. A None scope stands for changes made outside table resolvers, which + wait for nothing. Schedules all of them or, if it raises, none of them. Raises: - RuntimeError: if no sync is running, as the checkpoint could not be sent. + RuntimeError: if no sync is running, as the checkpoints could not be sent. """ with self._lock: if self._res is None: @@ -139,29 +145,37 @@ def schedule( "state checkpoints can only be flushed while the scheduler syncs, " "e.g. from a table resolver or the post-sync hook" ) - pending = _Pending(msg, keys) - for scope in {s for s in scopes if s is not None}: - # the scope's own resources are on the queue already: resolvers put each - # resource on it before they resume, so only its relations can be in flight - for child in scope.open_children: - child.waiters.append(pending) - pending.blockers += 1 - earlier = set() - for key in keys: - previous = self._unsent.get(key) - if previous is not None and previous not in earlier: - earlier.add(previous) - previous.dependents.append(pending) - pending.blockers += 1 - self._unsent[key] = pending - if pending.blockers == 0: - self._send([pending]) + for scopes, keys, msg in checkpoints: + self._schedule(scopes, keys, msg) + + def _schedule( + self, + scopes: Iterable[Optional[Scope]], + keys: List[str], + msg: SyncCheckpointMessage, + ) -> None: + pending = _Pending(msg, keys) + for scope in {s for s in scopes if s is not None}: + # the scope's own resources are on the queue already: resolvers put each + # resource on it before they resume, so only its relations can be in flight + for child in scope.open_children: + child.waiters.append(pending) + pending.blockers += 1 + earlier = set() + for key in keys: + previous = self._unsent.get(key) + if previous is not None and previous not in earlier: + earlier.add(previous) + previous.dependents.append(pending) + pending.blockers += 1 + self._unsent[key] = pending + if pending.blockers == 0: + self._send([pending]) def _send(self, ready: List[_Pending]) -> None: """Puts ready checkpoints on the queue, then the ones that only waited for them. Called with the lock held.""" while ready: pending = ready.pop() - pending.sent = True if self._res is not None: self._res.put(pending.msg) for key in pending.keys: diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index 3e3cfe0..6ea9abe 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -9,7 +9,6 @@ SyncInsertMessage, SyncMigrateTableMessage, SyncErrorMessage, - SyncCheckpointMessage, ) from cloudquery.sdk.internal.checkpoint import Checkpoints, Scope from cloudquery.sdk.schema import Resource @@ -35,6 +34,11 @@ def __init__(self) -> None: pass +class TableResolversSubmitted: + def __init__(self) -> None: + pass + + class Scheduler: def __init__( self, concurrency: int, queue_size: int = 0, max_depth: int = 3, logger=None @@ -95,8 +99,7 @@ def resolve_table( res: queue.Queue, scope: Optional[Scope] = None, ): - if scope is None: - scope = self.checkpoints.new_scope(None) + scope = scope or self.checkpoints.new_scope(None) try: with self.checkpoints.entered(scope): self._resolve_table(resolver, depth, client, parent_item, res, scope) @@ -197,19 +200,16 @@ def _sync( res: queue.Queue, deterministic_cq_id=False, ): - for resolver in resolvers: - clients = resolver.multiplex(client) - for client in clients: - res.put(TableResolverStarted()) - self._pools[0].submit( - self.resolve_table, - resolver, - 0, - client, - None, - res, - self.checkpoints.new_scope(None), - ) + try: + for resolver in resolvers: + clients = resolver.multiplex(client) + for client in clients: + res.put(TableResolverStarted()) + self._pools[0].submit( + self.resolve_table, resolver, 0, client, None, res + ) + finally: + res.put(TableResolversSubmitted()) def sync( self, client, resolvers: List[TableResolver], deterministic_cq_id=False @@ -221,31 +221,24 @@ def sync( try: thread = futures.ThreadPoolExecutor() thread.submit(self._sync, client, resolvers, res, deterministic_cq_id) + submitted = False total_table_resolvers = 0 finished_table_resolvers = 0 - while True: + while not (submitted and total_table_resolvers == finished_table_resolvers): message = res.get() - if isinstance(message, TableResolverStarted): + if isinstance(message, TableResolversSubmitted): + submitted = True + elif isinstance(message, TableResolverStarted): total_table_resolvers += 1 - if total_table_resolvers == finished_table_resolvers: - break - continue - if isinstance(message, TableResolverFinished): + elif isinstance(message, TableResolverFinished): finished_table_resolvers += 1 - if total_table_resolvers == finished_table_resolvers: - break - continue - yield message + else: + yield message self._post_sync_hook() # checkpoints the post-sync hook flushed, after every resource - while True: - try: - message = res.get_nowait() - except queue.Empty: - break - if isinstance(message, SyncCheckpointMessage): - yield message + while not res.empty(): + yield res.get_nowait() thread.shutdown(wait=True) finally: self.checkpoints.stop() diff --git a/cloudquery/sdk/stateclient/keys.py b/cloudquery/sdk/stateclient/keys.py index 0ec31c2..7e14c6c 100644 --- a/cloudquery/sdk/stateclient/keys.py +++ b/cloudquery/sdk/stateclient/keys.py @@ -18,7 +18,7 @@ def _escape(part: str) -> str: def state_key( table: str, *, - client_id: Optional[str] = None, + client_id: Optional[Union[str, int]] = None, **partition: Union[str, int], ) -> str: """ @@ -38,7 +38,7 @@ def state_key( """ parts = [_escape(table)] if client_id is not None: - parts.append(f"client_id{_ASSIGN}{_escape(client_id)}") + parts.append(f"client_id{_ASSIGN}{_escape(str(client_id))}") for name in sorted(partition): parts.append(f"{_escape(name)}{_ASSIGN}{_escape(str(partition[name]))}") return _SEPARATOR.join(parts) diff --git a/cloudquery/sdk/stateclient/stateclient.py b/cloudquery/sdk/stateclient/stateclient.py index 8a6b373..dfa245a 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/cloudquery/sdk/stateclient/stateclient.py @@ -120,9 +120,18 @@ def _send_checkpoints( groups: Dict[FrozenSet[Optional[Scope]], List[str]] = {} for key, scopes in changes.items(): groups.setdefault(frozenset(scopes), []).append(key) - for scopes, keys in groups.items(): - record = self.table.record([(key, values[key]) for key in keys]) - self.checkpoints.schedule(scopes, keys, SyncCheckpointMessage(record)) + self.checkpoints.schedule( + [ + ( + scopes, + keys, + SyncCheckpointMessage( + self.table.record([(key, values[key]) for key in keys]) + ), + ) + for scopes, keys in groups.items() + ] + ) @abstractmethod def migrate_state_table(self): diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py index 855f891..c608a87 100644 --- a/tests/scheduler/checkpoints.py +++ b/tests/scheduler/checkpoints.py @@ -1,6 +1,8 @@ +import contextvars import random import threading import time +from concurrent import futures from typing import Any, Dict, Generator, List, Optional, Tuple import pyarrow as pa @@ -300,3 +302,53 @@ def test_builder_uses_checkpoints_only_when_the_cli_supports_them(monkeypatch): assert supported.checkpoints is scheduler.checkpoints assert unsupported.checkpoints is None assert not_requested.checkpoints is None + + +class SingleItemResolver(TableResolver): + def __init__(self, name: str, multiplex_delay: float = 0): + super().__init__(Table(name, [Column("id", pa.int64())])) + self._multiplex_delay = multiplex_delay + + def multiplex(self, client): + time.sleep(self._multiplex_delay) + return [client] + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + yield {"id": 1} + + +def test_sync_waits_for_resolvers_submitted_after_earlier_ones_finished(): + scheduler = Scheduler(1) + resolvers = [ + SingleItemResolver("fast"), + SingleItemResolver("slow_to_multiplex", multiplex_delay=0.2), + ] + try: + messages = list(scheduler.sync(TestClient(), resolvers)) + finally: + scheduler.shutdown() + inserts = [m for m in messages if isinstance(m, SyncInsertMessage)] + assert len(inserts) == len(resolvers) + + +def test_scope_follows_copied_context_into_other_threads(): + scheduler = Scheduler(1) + scopes = [] + + class Resolver(TableResolver): + def __init__(self): + super().__init__(Table("items", [Column("id", pa.int64())])) + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + with futures.ThreadPoolExecutor(1) as pool: + ctx = contextvars.copy_context() + scopes.append( + pool.submit(ctx.run, scheduler.checkpoints.current_scope).result() + ) + yield {"id": 1} + + try: + list(scheduler.sync(TestClient(), [Resolver()])) + finally: + scheduler.shutdown() + assert scopes[0] is not None diff --git a/tests/stateclient/keys.py b/tests/stateclient/keys.py index a29f335..d431b08 100644 --- a/tests/stateclient/keys.py +++ b/tests/stateclient/keys.py @@ -31,3 +31,7 @@ def test_tables_clients_and_partitions_never_share_a_key(): state_key("t", a=1, b=2), ] assert len(set(keys)) == len(keys) + + +def test_client_id_that_is_not_a_string(): + assert state_key("t", client_id=123) == "t|client_id=123" From 847a92bdf7dd73bcf63401fa9e3796b20a9ebf8a Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 14:46:40 +0300 Subject: [PATCH 2/2] fix: Reject checkpoints from resolvers of an ended sync A Scheduler reused across syncs kept resolvers of a cancelled sync running. Their flushes went onto the next sync's stream, so the CLI could persist cursors past the rows it committed. --- cloudquery/sdk/internal/checkpoint.py | 21 +++++++++---- tests/scheduler/checkpoints.py | 43 +++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py index df038f0..394d6a3 100644 --- a/cloudquery/sdk/internal/checkpoint.py +++ b/cloudquery/sdk/internal/checkpoint.py @@ -30,8 +30,10 @@ class Scope: spawned is finished. """ - def __init__(self, parent: Optional[Scope]): + def __init__(self, parent: Optional[Scope], sync: Optional[object]): self.parent = parent + # the sync the scope belongs to, so that resolvers still running after it ended cannot checkpoint into the next one + self.sync = sync # child scopes that are not finished yet self.open_children: Set[Scope] = set() self.resolving = True @@ -45,9 +47,10 @@ def finished(self) -> bool: class _Pending: """A checkpoint waiting for scopes to finish and for earlier checkpoints of its keys to be sent.""" - def __init__(self, msg: SyncCheckpointMessage, keys: List[str]): + def __init__(self, msg: SyncCheckpointMessage, keys: List[str], sync: object): self.msg = msg self.keys = keys + self.sync = sync self.blockers = 0 self.dependents: List[_Pending] = [] @@ -68,22 +71,25 @@ def __init__(self): self._res: Optional[queue.Queue] = None # newest unsent checkpoint of each key self._unsent: Dict[str, _Pending] = {} + self._sync: Optional[object] = None def start(self, res: queue.Queue) -> None: """Sends later checkpoints on res. Called by the scheduler when a sync starts.""" with self._lock: self._res = res + self._sync = object() self._unsent = {} def stop(self) -> None: """Drops unsent checkpoints and rejects later ones. Called by the scheduler when a sync ends.""" with self._lock: self._res = None + self._sync = None self._unsent = {} def new_scope(self, parent: Optional[Scope]) -> Scope: """Returns a scope for a table resolver call, before the call is submitted.""" - scope = Scope(parent) + scope = Scope(parent, parent.sync if parent is not None else self._sync) if parent is not None: with self._lock: parent.open_children.add(scope) @@ -145,6 +151,11 @@ def schedule( "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( + "state checkpoints of a table resolver can only be flushed during its own sync" + ) for scopes, keys, msg in checkpoints: self._schedule(scopes, keys, msg) @@ -154,7 +165,7 @@ def _schedule( keys: List[str], msg: SyncCheckpointMessage, ) -> None: - pending = _Pending(msg, keys) + pending = _Pending(msg, keys, self._sync) for scope in {s for s in scopes if s is not None}: # the scope's own resources are on the queue already: resolvers put each # resource on it before they resume, so only its relations can be in flight @@ -176,7 +187,7 @@ def _send(self, ready: List[_Pending]) -> None: """Puts ready checkpoints on the queue, then the ones that only waited for them. Called with the lock held.""" while ready: pending = ready.pop() - if self._res is not None: + if pending.sync is self._sync: self._res.put(pending.msg) for key in pending.keys: if self._unsent.get(key) is pending: diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py index c608a87..61fc182 100644 --- a/tests/scheduler/checkpoints.py +++ b/tests/scheduler/checkpoints.py @@ -352,3 +352,46 @@ def resolve(self, client, parent_resource) -> Generator[Any, None, None]: finally: scheduler.shutdown() assert scopes[0] is not None + + +class CursorResolver(TableResolver): + def __init__(self, state: StateClient, tag: str): + super().__init__(Table("items", [Column("id", pa.int64())])) + self._state = state + self._tag = tag + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + for item_id in range(1, 30): + time.sleep(0.02) + yield {"id": item_id} + self._state.set_key("cursor", f"{self._tag}-{item_id}") + self._state.flush() + + +def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one(): + scheduler = Scheduler(4) + try: + ended = scheduler.sync( + TestClient(), + [ + CursorResolver( + MemoryStateClient(checkpoints=scheduler.checkpoints), "ended" + ) + ], + ) + for _ in range(5): + next(ended) + ended.close() + + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + scheduler.set_post_sync_hook(state.flush) + messages = list(scheduler.sync(TestClient(), [CursorResolver(state, "next")])) + finally: + scheduler.shutdown() + cursors = [ + checkpoint_rows(m)["cursor"] + for m in messages + if isinstance(m, SyncCheckpointMessage) + ] + assert cursors + assert all(cursor.startswith("next-") for cursor in cursors)