From 0dc34a5a7976c2b775d02d687bdbe88d1326bbf3 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 19 Oct 2026 14:04:02 -0400 Subject: [PATCH 1/6] feat: Send state as checkpoints after the resources it covers When the CLI sets Sync.Request.supports_checkpoints, a state client built with the scheduler's checkpoints sends flushed state as Sync.MessageCheckpoint instead of writing it to the backend. The CLI then persists it only after every destination commits the records synced before it, so a failed sync keeps the progress it made without moving a cursor past records that were never written. Each checkpoint goes on the sync stream after the resources the table resolvers that changed its keys produced before the flush, including the relations spawned from them, and after earlier unsent checkpoints of the same keys. Also makes the state client thread-safe, clears changed keys after a successful flush and keeps them after a failed one. --- cloudquery/sdk/internal/checkpoint.py | 174 ++++++++++ .../sdk/internal/servers/plugin_v3/plugin.py | 24 +- cloudquery/sdk/message/__init__.py | 1 + cloudquery/sdk/message/sync.py | 10 + cloudquery/sdk/plugin/plugin.py | 2 + cloudquery/sdk/scheduler/scheduler.py | 87 +++-- cloudquery/sdk/stateclient/stateclient.py | 108 +++++-- pyproject.toml | 2 +- tests/scheduler/checkpoints.py | 302 ++++++++++++++++++ tests/serve/checkpoints.py | 50 +++ 10 files changed, 711 insertions(+), 49 deletions(-) create mode 100644 cloudquery/sdk/internal/checkpoint.py create mode 100644 tests/scheduler/checkpoints.py create mode 100644 tests/serve/checkpoints.py diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py new file mode 100644 index 0000000..1b04b0e --- /dev/null +++ b/cloudquery/sdk/internal/checkpoint.py @@ -0,0 +1,174 @@ +""" +Orders state checkpoints after the resources they summarize. + +When the CLI supports checkpoints, it persists the state rows a source sends as +``Sync.MessageCheckpoint`` once every destination has committed the messages sent +before them. A checkpoint must therefore go on the sync stream only after every +resource its state covers: the resources produced, before the flush, by the table +resolvers that changed its keys, and the relations those resources spawned. + +The scheduler sends every message through a single queue, so a checkpoint put on +the queue once those resources are on it is ordered after them on the stream. +""" + +from __future__ import annotations + +import queue +import threading +from contextlib import contextmanager +from typing import Dict, Iterable, Iterator, List, Optional, Set + +from cloudquery.sdk.message import SyncCheckpointMessage + + +class Scope: + """ + A table resolver call and the relations its resources spawn. + + A scope is finished once its resolver returned and every child scope it + spawned is finished. + """ + + def __init__(self, parent: Optional[Scope]): + self.parent = parent + # child scopes that are not finished yet + self.open_children: Set[Scope] = set() + self.resolving = True + # checkpoints waiting for this scope to finish + self.waiters: List[_Pending] = [] + + def finished(self) -> bool: + return not self.resolving and not self.open_children + + +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]): + self.msg = msg + self.keys = keys + self.blockers = 0 + self.dependents: List[_Pending] = [] + self.sent = False + + +class Checkpoints: + """ + Sends state checkpoints on the scheduler's queue after the resources they cover. + + The scheduler owns an instance and passes it to the state client through + :meth:`cloudquery.sdk.stateclient.StateClientBuilder.build`. + """ + + def __init__(self): + self._lock = threading.Lock() + self._local = threading.local() + self._res: Optional[queue.Queue] = None + # newest unsent checkpoint of each key + self._unsent: Dict[str, _Pending] = {} + + 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._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._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) + if parent is not None: + with self._lock: + parent.open_children.add(scope) + return 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 + try: + yield + finally: + self._local.scope = previous + + 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) + + def resolver_finished(self, scope: Scope) -> None: + """ + Marks the resolver of scope as returned, once every resource it produced is on the queue. + """ + with self._lock: + scope.resolving = False + ready: List[_Pending] = [] + while scope is not None and scope.finished(): + for pending in scope.waiters: + pending.blockers -= 1 + if pending.blockers == 0: + ready.append(pending) + scope.waiters = [] + if scope.parent is not None: + scope.parent.open_children.discard(scope) + scope = scope.parent + self._send(ready) + + def schedule( + self, + scopes: Iterable[Optional[Scope]], + keys: List[str], + msg: 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. + + Raises: + RuntimeError: if no sync is running, as the checkpoint could not be sent. + """ + with self._lock: + if self._res is None: + raise RuntimeError( + "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]) + + 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: + if self._unsent.get(key) is pending: + del self._unsent[key] + for dependent in pending.dependents: + dependent.blockers -= 1 + if dependent.blockers == 0: + ready.append(dependent) + pending.dependents = [] diff --git a/cloudquery/sdk/internal/servers/plugin_v3/plugin.py b/cloudquery/sdk/internal/servers/plugin_v3/plugin.py index 6531e98..0dd6a1f 100644 --- a/cloudquery/sdk/internal/servers/plugin_v3/plugin.py +++ b/cloudquery/sdk/internal/servers/plugin_v3/plugin.py @@ -8,12 +8,18 @@ SyncInsertMessage, SyncMigrateTableMessage, SyncErrorMessage, + SyncCheckpointMessage, WriteInsertMessage, WriteMigrateTableMessage, WriteMessage, WriteDeleteStale, ) -from cloudquery.sdk.plugin.plugin import Plugin, SyncOptions, TableOptions +from cloudquery.sdk.plugin.plugin import ( + BackendOptions, + Plugin, + SyncOptions, + TableOptions, +) from cloudquery.sdk.schema import tables_to_arrow_schemas, Table from cloudquery.sdk.schema.table import flatten_tables @@ -64,7 +70,11 @@ def Sync(self, request, context): skip_dependent_tables=request.skip_dependent_tables, skip_tables=request.skip_tables, tables=request.tables, - backend_options=request.backend, + backend_options=BackendOptions( + connection=request.backend.connection, + table_name=request.backend.table_name, + supports_checkpoints=request.supports_checkpoints, + ), ) for msg in self._plugin.sync(options): @@ -78,6 +88,16 @@ def Sync(self, request, context): yield plugin_pb2.Sync.Response( migrate_table=plugin_pb2.Sync.MessageMigrateTable(table=buf) ) + elif isinstance(msg, SyncCheckpointMessage): + if not request.supports_checkpoints: + # the state would be lost, as the CLI does not persist checkpoints + raise RuntimeError( + "plugin sent a state checkpoint, but the CLI does not support checkpoints" + ) + buf = arrow.record_to_bytes(msg.record) + yield plugin_pb2.Sync.Response( + checkpoint=plugin_pb2.Sync.MessageCheckpoint(record=buf) + ) elif isinstance(msg, SyncErrorMessage) and request.withErrorMessages: yield plugin_pb2.Sync.Response( error=plugin_pb2.Sync.MessageError( diff --git a/cloudquery/sdk/message/__init__.py b/cloudquery/sdk/message/__init__.py index 25592ef..e52698f 100644 --- a/cloudquery/sdk/message/__init__.py +++ b/cloudquery/sdk/message/__init__.py @@ -3,6 +3,7 @@ SyncInsertMessage, SyncMigrateTableMessage, SyncErrorMessage, + SyncCheckpointMessage, ) from .write import ( WriteMessage, diff --git a/cloudquery/sdk/message/sync.py b/cloudquery/sdk/message/sync.py index f971a12..7e1538a 100644 --- a/cloudquery/sdk/message/sync.py +++ b/cloudquery/sdk/message/sync.py @@ -19,3 +19,13 @@ class SyncErrorMessage(SyncMessage): def __init__(self, table_name: str, error: str): self.table_name = table_name self.error = error + + +class SyncCheckpointMessage(SyncMessage): + """ + State backend rows for the CLI to persist once every message sent before this one + is committed to all destinations. Only sent when the CLI supports checkpoints. + """ + + def __init__(self, record: pa.RecordBatch): + self.record = record diff --git a/cloudquery/sdk/plugin/plugin.py b/cloudquery/sdk/plugin/plugin.py index d03dbef..9d93f7d 100644 --- a/cloudquery/sdk/plugin/plugin.py +++ b/cloudquery/sdk/plugin/plugin.py @@ -18,6 +18,8 @@ class TableOptions: class BackendOptions: connection: str = None table_name: str = None + # set when the CLI persists state sent as checkpoints on the sync stream + supports_checkpoints: bool = False @dataclass diff --git a/cloudquery/sdk/scheduler/scheduler.py b/cloudquery/sdk/scheduler/scheduler.py index bc00b99..3e3cfe0 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -9,7 +9,9 @@ SyncInsertMessage, SyncMigrateTableMessage, SyncErrorMessage, + SyncCheckpointMessage, ) +from cloudquery.sdk.internal.checkpoint import Checkpoints, Scope from cloudquery.sdk.schema import Resource from cloudquery.sdk.stateclient.stateclient import StateClient from .table_resolver import TableResolver, Client @@ -39,6 +41,8 @@ def __init__( ): self._post_sync_hook = lambda: None self._queue = queue.Queue() + # pass to StateClientBuilder.build so that state flushes are sent as checkpoints, when the CLI supports them + self.checkpoints = Checkpoints() self._max_depth = max_depth if logger is None: self._logger = structlog.get_logger() @@ -89,7 +93,28 @@ def resolve_table( client: Client, parent_item: Optional[Resource], res: queue.Queue, + scope: Optional[Scope] = None, ): + if scope is None: + scope = self.checkpoints.new_scope(None) + try: + with self.checkpoints.entered(scope): + self._resolve_table(resolver, depth, client, parent_item, res, scope) + finally: + # after every resource of the resolver is on the queue, so its checkpoints follow them + self.checkpoints.resolver_finished(scope) + res.put(TableResolverFinished()) + + def _resolve_table( + self, + resolver: TableResolver, + depth: int, + client: Client, + parent_item: Optional[Resource], + res: queue.Queue, + scope: Scope, + ): + total_resources = 0 try: if depth == 0: self._logger.info( @@ -105,7 +130,6 @@ def resolve_table( table=resolver.table.name, depth=depth, ) - total_resources = 0 for item in resolver.resolve(client, parent_item): try: resource = self.resolve_resource( @@ -136,6 +160,7 @@ def resolve_table( client, resource, res, + self.checkpoints.new_scope(scope), ) total_resources += 1 if depth == 0: @@ -164,8 +189,6 @@ def resolve_table( exc_info=e, ) res.put(SyncErrorMessage(resolver.table.name, str(e))) - finally: - res.put(TableResolverFinished()) def _sync( self, @@ -179,7 +202,13 @@ def _sync( for client in clients: res.put(TableResolverStarted()) self._pools[0].submit( - self.resolve_table, resolver, 0, client, None, res + self.resolve_table, + resolver, + 0, + client, + None, + res, + self.checkpoints.new_scope(None), ) def sync( @@ -188,26 +217,38 @@ def sync( res = queue.Queue() yield from self._send_migrate_table_messages(resolvers) - thread = futures.ThreadPoolExecutor() - thread.submit(self._sync, client, resolvers, res, deterministic_cq_id) - total_table_resolvers = 0 - finished_table_resolvers = 0 - while True: - message = res.get() - if isinstance(message, TableResolverStarted): - total_table_resolvers += 1 - if total_table_resolvers == finished_table_resolvers: - break - continue - if isinstance(message, TableResolverFinished): - finished_table_resolvers += 1 - if total_table_resolvers == finished_table_resolvers: - break - continue - yield message + self.checkpoints.start(res) + try: + thread = futures.ThreadPoolExecutor() + thread.submit(self._sync, client, resolvers, res, deterministic_cq_id) + total_table_resolvers = 0 + finished_table_resolvers = 0 + while True: + message = res.get() + if isinstance(message, TableResolverStarted): + total_table_resolvers += 1 + if total_table_resolvers == finished_table_resolvers: + break + continue + if isinstance(message, TableResolverFinished): + finished_table_resolvers += 1 + if total_table_resolvers == finished_table_resolvers: + break + continue + yield message - self._post_sync_hook() - thread.shutdown(wait=True) + 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 + thread.shutdown(wait=True) + finally: + self.checkpoints.stop() def _send_migrate_table_messages( self, resolvers: List[TableResolver] diff --git a/cloudquery/sdk/stateclient/stateclient.py b/cloudquery/sdk/stateclient/stateclient.py index 0b0dded..08ffe73 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/cloudquery/sdk/stateclient/stateclient.py @@ -1,11 +1,14 @@ from __future__ import annotations import grpc +import threading from abc import ABC, abstractmethod -from typing import Generator, Optional, Tuple +from typing import Dict, FrozenSet, Generator, List, Optional, Set, Tuple import pyarrow as pa from cloudquery.sdk import schema +from cloudquery.sdk.internal.checkpoint import Checkpoints, Scope +from cloudquery.sdk.message import SyncCheckpointMessage from cloudquery.sdk.plugin.plugin import BackendOptions from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc, arrow from functools import wraps @@ -21,6 +24,13 @@ class StateClientBuilder: Args: backend_options (BackendOptions): which has `connection` & `table_name` strings. + checkpoints (Checkpoints): optional, the `checkpoints` of the `Scheduler` that syncs. + When the CLI supports checkpoints, `flush` then sends the state on the sync stream + instead of writing it to the backend, and the CLI persists it once the destinations + commit the resources synced before it. Changes made within a table resolver wait only + for the resources that resolver produced before the flush, and their relations. + Changes made outside table resolvers wait for nothing, so flush them after the + resources they summarize are synced, e.g. in the scheduler's post-sync hook. Returns: `ConnectedStateClient` or `NoOpStateClient` @@ -28,11 +38,17 @@ class StateClientBuilder: """ @staticmethod - def build(*, backend_options: BackendOptions) -> StateClient: + def build( + *, backend_options: BackendOptions, checkpoints: Optional[Checkpoints] = None + ) -> StateClient: if not backend_options or not backend_options.table_name: return NoOpStateClient(backend_options=backend_options) - return ConnectedStateClient(backend_options=backend_options) + if not getattr(backend_options, "supports_checkpoints", False): + checkpoints = None + return ConnectedStateClient( + backend_options=backend_options, checkpoints=checkpoints + ) class StateClient(ABC): @@ -43,31 +59,70 @@ class StateClient(ABC): so it is a succinct overview of what a state client does. """ - def __init__(self, *, backend_options: BackendOptions): + def __init__( + self, + *, + backend_options: BackendOptions, + checkpoints: Optional[Checkpoints] = None, + ): self.mem = {} - self.changes = {} + # changed keys, each with the scopes of the table resolvers that changed it since the + # last flush (None for changes made outside table resolvers) + self.changes: Dict[str, Set[Optional[Scope]]] = {} self.connection = getattr(backend_options, "connection", None) self.table = Table(getattr(backend_options, "table_name", None)) + self.checkpoints = checkpoints + # guards mem and changes, which table resolvers change from several threads + self._lock = threading.RLock() + # serializes flushes, so that a key's flushed values reach the state in the order they were set + self._flush_lock = threading.Lock() self.migrate_state_table() self.read_all_state() def get_key(self, key: str) -> Optional[str]: - return self.mem.get(key) + with self._lock: + return self.mem.get(key) def set_key(self, key: str, value: str) -> None: - self.mem[key] = value - self.changes[key] = True + scope = self.checkpoints.current_scope() if self.checkpoints else None + with self._lock: + self.mem[key] = value + self.changes.setdefault(key, set()).add(scope) def flush(self): - if not self.changes: - return - - self.write_all_state(self._changed_keys()) - - def _changed_keys(self): - for key, _ in self.changes.items(): - yield (key, self.mem[key]) + """ + 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. + """ + with self._flush_lock: + with self._lock: + if not self.changes: + return + changes, self.changes = self.changes, {} + values = {key: self.mem[key] for key in changes} + try: + if self.checkpoints is None: + self.write_all_state(iter(values.items())) + else: + self._send_checkpoints(changes, values) + except Exception: + with self._lock: + for key, scopes in changes.items(): + self.changes.setdefault(key, set()).update(scopes) + raise + + def _send_checkpoints( + self, changes: Dict[str, Set[Optional[Scope]]], values: Dict[str, str] + ): + # one checkpoint per set of scopes, so each waits only for the resources of the scopes that changed its keys + 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)) @abstractmethod def migrate_state_table(self): @@ -93,6 +148,9 @@ def get_key(self, key: str) -> Optional[str]: def set_key(self, key: str, value: str) -> None: pass + def flush(self): + pass + def migrate_state_table(self): pass @@ -144,13 +202,7 @@ def write_all_state( backend_plugin.Write(self._write_request(k, v) for k, v in changed_keys) def _write_request(self, key: str, value: str) -> plugin_pb2.Write.Request: - record = pa.RecordBatch.from_arrays( - [ - pa.array([key]), - pa.array([value]), - ], - schema=self.table.arrow_schema(), - ) + record = self.table.record([(key, value)]) return plugin_pb2.Write.Request( insert=plugin_pb2.Write.MessageInsert(record=arrow.record_to_bytes(record)) ) @@ -197,6 +249,16 @@ def arrow_schema(self): def bytes(self): return self._bytes + def record(self, rows: List[Tuple[str, str]]) -> pa.RecordBatch: + """Returns the state rows as a record of this table, with the table name in its schema metadata.""" + return pa.RecordBatch.from_arrays( + [ + pa.array([key for key, _ in rows], type=pa.string()), + pa.array([value for _, value in rows], type=pa.string()), + ], + schema=self._arrow_schema, + ) + def _state_table_schema(self): return schema.Table( name=self.name, diff --git a/pyproject.toml b/pyproject.toml index a499511..13d7df4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Topic :: Internet", ] dependencies = [ - "cloudquery-plugin-pb>=0.0.54", + "cloudquery-plugin-pb>=0.0.55", "grpcio>=1.78.0", "grpcio-tools>=1.78.0", "Jinja2>=3.1.6", diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py new file mode 100644 index 0000000..855f891 --- /dev/null +++ b/tests/scheduler/checkpoints.py @@ -0,0 +1,302 @@ +import random +import threading +import time +from typing import Any, Dict, Generator, List, Optional, Tuple + +import pyarrow as pa +import pytest + +from cloudquery.sdk.message import ( + SyncCheckpointMessage, + SyncInsertMessage, + SyncMessage, +) +from cloudquery.sdk.plugin.plugin import BackendOptions +from cloudquery.sdk.scheduler import Scheduler, TableResolver +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.stateclient import ( + ConnectedStateClient, + StateClient, + StateClientBuilder, +) + +STATE_TABLE = "cq_state_test" + + +class MemoryStateClient(StateClient): + """A state client whose backend is a dict.""" + + def __init__(self, checkpoints=None, fail_writes=False): + self.backend: Dict[str, str] = {} + self.writes: List[List[Tuple[str, str]]] = [] + self.fail_writes = fail_writes + super().__init__( + backend_options=BackendOptions( + connection="unused", table_name=STATE_TABLE, supports_checkpoints=True + ), + checkpoints=checkpoints, + ) + + def migrate_state_table(self): + pass + + def read_all_state(self): + self.mem.update(self.backend) + + def write_all_state(self, changed_keys): + rows = list(changed_keys) + if self.fail_writes: + raise ConnectionError("backend unavailable") + self.writes.append(rows) + self.backend.update(rows) + + +class TestClient: + def id(self): + return "test_client" + + +def checkpoint_rows(msg: SyncCheckpointMessage) -> Dict[str, str]: + rows = msg.record.to_pylist() + return {row["key"]: row["value"] for row in rows} + + +def sleep_a_little(): + time.sleep(random.uniform(0, 0.005)) + + +class ChildResolver(TableResolver): + """Yields children slowly, so they finish after the parent moved on.""" + + def __init__(self, children: int, state: Optional[StateClient] = None): + super().__init__( + Table( + "child", + [Column("parent_id", pa.int64()), Column("child_id", pa.int64())], + ) + ) + self._children = children + self._state = state + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + parent_id = parent_resource.item["id"] + for child_id in range(self._children): + sleep_a_little() + yield {"parent_id": parent_id, "child_id": child_id} + if self._state is not None: + self._state.set_key(f"child-{parent_id}", str(self._children - 1)) + self._state.flush() + + +class ParentResolver(TableResolver): + """Yields parents and checkpoints its cursor after each one, as an incremental table would.""" + + def __init__(self, state: StateClient, parents: int, children: int, child_state): + super().__init__( + Table("parent", [Column("id", pa.int64())]), + child_resolvers=[ChildResolver(children, child_state)], + ) + self._state = state + self._parents = parents + + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: + for parent_id in range(self._parents): + yield {"id": parent_id} + self._state.set_key("parent-cursor", str(parent_id)) + self._state.flush() + + +def run_sync( + parents=20, children=5, child_checkpoints=False, concurrency=8 +) -> Tuple[List[SyncMessage], MemoryStateClient]: + scheduler = Scheduler(concurrency) + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + scheduler.set_post_sync_hook(state.flush) + resolver = ParentResolver( + state, parents, children, state if child_checkpoints else None + ) + try: + messages = list(scheduler.sync(TestClient(), [resolver])) + finally: + scheduler.shutdown() + return messages, state + + +def test_parent_checkpoints_follow_the_relations_of_the_resources_they_cover(): + parents, children = 20, 5 + messages, state = run_sync(parents, children) + + assert state.writes == [], "state must not be written to the backend directly" + seen_children = set() + cursors = [] + for msg in messages: + if isinstance(msg, SyncInsertMessage) and msg.record.schema.metadata[ + METADATA_TABLE_NAME + ] == (b"child"): + row = msg.record.to_pylist()[0] + seen_children.add((row["parent_id"], row["child_id"])) + elif isinstance(msg, SyncCheckpointMessage): + cursor = int(checkpoint_rows(msg)["parent-cursor"]) + cursors.append(cursor) + for parent_id in range(cursor + 1): + for child_id in range(children): + assert (parent_id, child_id) in seen_children, ( + f"checkpoint of cursor {cursor} sent before child " + f"{child_id} of parent {parent_id}" + ) + assert cursors, "no checkpoints were sent" + assert cursors == sorted(cursors), "a key's checkpoints must keep their order" + assert cursors[-1] == parents - 1 + + +def test_child_checkpoints_follow_their_own_resources(): + parents, children = 10, 4 + messages, _ = run_sync(parents, children, child_checkpoints=True) + + seen_children = set() + child_keys = set() + for msg in messages: + if isinstance(msg, SyncInsertMessage) and msg.record.schema.metadata[ + METADATA_TABLE_NAME + ] == (b"child"): + row = msg.record.to_pylist()[0] + seen_children.add((row["parent_id"], row["child_id"])) + elif isinstance(msg, SyncCheckpointMessage): + for key in checkpoint_rows(msg): + if key.startswith("child-"): + parent_id = int(key.split("-")[1]) + child_keys.add(key) + for child_id in range(children): + assert (parent_id, child_id) in seen_children + assert len(child_keys) == parents + + +def test_checkpoint_records_are_state_table_rows(): + messages, _ = run_sync(parents=2, children=1) + checkpoint = next(m for m in messages if isinstance(m, SyncCheckpointMessage)) + schema = checkpoint.record.schema + assert schema.metadata[METADATA_TABLE_NAME] == STATE_TABLE.encode() + assert schema.field("key").type == pa.string() + assert schema.field("value").type == pa.string() + + +def test_post_sync_hook_flush_is_sent_after_every_resource(): + scheduler = Scheduler(2) + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + + class Resolver(TableResolver): + def __init__(self): + super().__init__(Table("items", [Column("id", pa.int64())])) + + def resolve(self, client, parent_resource): + yield {"id": 1} + yield {"id": 2} + + def hook(): + # outside table resolvers, so it waits for nothing: the hook runs after every resource + state.set_key("sync-finished", "yes") + state.flush() + + scheduler.set_post_sync_hook(hook) + messages = list(scheduler.sync(TestClient(), [Resolver()])) + scheduler.shutdown() + + assert isinstance(messages[-1], SyncCheckpointMessage) + assert checkpoint_rows(messages[-1]) == {"sync-finished": "yes"} + assert sum(isinstance(m, SyncInsertMessage) for m in messages) == 2 + + +def test_flush_writes_to_the_backend_without_checkpoints(): + state = MemoryStateClient(checkpoints=None) + state.set_key("a", "1") + state.set_key("b", "2") + state.flush() + state.flush() # nothing changed since + + assert state.writes == [[("a", "1"), ("b", "2")]] + state.set_key("a", "3") + state.flush() + assert state.writes[-1] == [("a", "3")] + + +def test_failed_flush_keeps_the_changes(): + state = MemoryStateClient(checkpoints=None, fail_writes=True) + state.set_key("a", "1") + with pytest.raises(ConnectionError): + state.flush() + + state.fail_writes = False + state.flush() + assert state.backend == {"a": "1"} + + +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): + state.flush() + assert "a" in state.changes + assert state.writes == [] + scheduler.shutdown() + + +def test_concurrent_flushes_send_the_latest_value_last(): + scheduler = Scheduler(16) + state = MemoryStateClient(checkpoints=scheduler.checkpoints) + counter = iter(range(10**6)) + lock = threading.Lock() + + class Resolver(TableResolver): + def __init__(self, name): + super().__init__(Table(name, [Column("id", pa.int64())])) + + def resolve(self, client, parent_resource): + for i in range(50): + yield {"id": i} + with lock: + value = next(counter) + state.set_key("shared", f"{value:08d}") + state.flush() + + messages = list( + scheduler.sync(TestClient(), [Resolver(f"items_{i}") for i in range(8)]) + ) + scheduler.shutdown() + + values = [ + checkpoint_rows(m)["shared"] + for m in messages + if isinstance(m, SyncCheckpointMessage) + ] + assert values == sorted(values) + assert values[-1] == f"{8 * 50 - 1:08d}" + + +def test_builder_uses_checkpoints_only_when_the_cli_supports_them(monkeypatch): + for method in ("migrate_state_table", "read_all_state", "write_all_state"): + monkeypatch.setattr(ConnectedStateClient, method, lambda self, *a: None) + scheduler = Scheduler(1) + + supported = StateClientBuilder.build( + backend_options=BackendOptions( + connection="localhost:1", table_name="t", supports_checkpoints=True + ), + checkpoints=scheduler.checkpoints, + ) + unsupported = StateClientBuilder.build( + backend_options=BackendOptions(connection="localhost:1", table_name="t"), + checkpoints=scheduler.checkpoints, + ) + not_requested = StateClientBuilder.build( + backend_options=BackendOptions( + connection="localhost:1", table_name="t", supports_checkpoints=True + ), + ) + scheduler.shutdown() + + assert supported.checkpoints is scheduler.checkpoints + assert unsupported.checkpoints is None + assert not_requested.checkpoints is None diff --git a/tests/serve/checkpoints.py b/tests/serve/checkpoints.py new file mode 100644 index 0000000..3b1df44 --- /dev/null +++ b/tests/serve/checkpoints.py @@ -0,0 +1,50 @@ +import pyarrow as pa +import pytest +from cloudquery.plugin_v3 import arrow, plugin_pb2 + +from cloudquery.sdk.internal.servers.plugin_v3.plugin import PluginServicer +from cloudquery.sdk.message import SyncCheckpointMessage, SyncInsertMessage +from cloudquery.sdk.plugin.plugin import Plugin, SyncOptions +from cloudquery.sdk.stateclient.stateclient import Table as StateTable + +state_table = StateTable("cq_state_test") + + +class CheckpointingPlugin(Plugin): + def __init__(self): + super().__init__("checkpointing", "development") + self.options = None + + def sync(self, options: SyncOptions): + self.options = options + yield SyncInsertMessage(pa.record_batch([[1]], names=["id"])) + yield SyncCheckpointMessage(state_table.record([("cursor", "1")])) + + +def sync_request(supports_checkpoints: bool): + return plugin_pb2.Sync.Request( + tables=["*"], + backend=plugin_pb2.Sync.BackendOptions( + table_name="cq_state_test", connection="localhost:7777" + ), + supports_checkpoints=supports_checkpoints, + ) + + +def test_sync_sends_checkpoints_and_passes_the_cli_support(): + plugin = CheckpointingPlugin() + responses = list(PluginServicer(plugin).Sync(sync_request(True), None)) + + assert plugin.options.backend_options.supports_checkpoints + assert plugin.options.backend_options.table_name == "cq_state_test" + assert plugin.options.backend_options.connection == "localhost:7777" + assert [r.WhichOneof("message") for r in responses] == ["insert", "checkpoint"] + record = arrow.new_record_from_bytes(responses[1].checkpoint.record) + assert record.to_pylist() == [{"key": "cursor", "value": "1"}] + + +def test_sync_refuses_checkpoints_the_cli_would_drop(): + plugin = CheckpointingPlugin() + with pytest.raises(RuntimeError): + list(PluginServicer(plugin).Sync(sync_request(False), None)) + assert not plugin.options.backend_options.supports_checkpoints From 8aff4b7bb10a7b6cf48f13fc5ef6d3b016bde251 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 19 Oct 2026 14:47:12 -0400 Subject: [PATCH 2/6] Read supports_checkpoints directly, as BackendOptions defaults it --- cloudquery/sdk/stateclient/stateclient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cloudquery/sdk/stateclient/stateclient.py b/cloudquery/sdk/stateclient/stateclient.py index 08ffe73..8a6b373 100644 --- a/cloudquery/sdk/stateclient/stateclient.py +++ b/cloudquery/sdk/stateclient/stateclient.py @@ -44,7 +44,7 @@ def build( if not backend_options or not backend_options.table_name: return NoOpStateClient(backend_options=backend_options) - if not getattr(backend_options, "supports_checkpoints", False): + if not backend_options.supports_checkpoints: checkpoints = None return ConnectedStateClient( backend_options=backend_options, checkpoints=checkpoints From 9eecf8fb1b21d4dd30a8c9cf02679834e711932d Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 19 Oct 2026 15:30:57 -0400 Subject: [PATCH 3/6] feat: Add state_key to build per-table and per-partition state keys --- cloudquery/sdk/stateclient/__init__.py | 1 + cloudquery/sdk/stateclient/keys.py | 44 ++++++++++++++++++++++++++ tests/stateclient/__init__.py | 0 tests/stateclient/keys.py | 33 +++++++++++++++++++ 4 files changed, 78 insertions(+) create mode 100644 cloudquery/sdk/stateclient/keys.py create mode 100644 tests/stateclient/__init__.py create mode 100644 tests/stateclient/keys.py diff --git a/cloudquery/sdk/stateclient/__init__.py b/cloudquery/sdk/stateclient/__init__.py index e69de29..ddff7be 100644 --- a/cloudquery/sdk/stateclient/__init__.py +++ b/cloudquery/sdk/stateclient/__init__.py @@ -0,0 +1 @@ +from .keys import state_key diff --git a/cloudquery/sdk/stateclient/keys.py b/cloudquery/sdk/stateclient/keys.py new file mode 100644 index 0000000..0ec31c2 --- /dev/null +++ b/cloudquery/sdk/stateclient/keys.py @@ -0,0 +1,44 @@ +"""Builds state keys, so that plugins choose how finely they track progress without inventing key formats.""" + +from typing import Optional, Union + +_SEPARATOR = "|" +_ASSIGN = "=" +_ESCAPE = "\\" + + +def _escape(part: str) -> str: + return ( + part.replace(_ESCAPE, _ESCAPE * 2) + .replace(_SEPARATOR, _ESCAPE + _SEPARATOR) + .replace(_ASSIGN, _ESCAPE + _ASSIGN) + ) + + +def state_key( + table: str, + *, + client_id: Optional[str] = None, + **partition: Union[str, int], +) -> str: + """ + Returns the state key of a table's progress, optionally for one client and one partition of the table. + + A table with one cursor uses ``state_key("items")``. A table that tracks progress per parent, account or + other partition passes the fields that identify it, e.g. ``state_key("responses", form_id=form["id"])``. + Each such key is its own state row: with checkpoints, a partition's state waits only for the resources of + the table resolver that set it, and the CLI persists the keys changed since its last commit together. + + Keys are readable and deterministic: partition fields are sorted, and separators in names and values are + escaped, so different tables, clients or partitions never share a key. + + Example: + >>> state_key("responses", client_id="acct-1", form_id=123) + 'responses|client_id=acct-1|form_id=123' + """ + parts = [_escape(table)] + if client_id is not None: + parts.append(f"client_id{_ASSIGN}{_escape(client_id)}") + for name in sorted(partition): + parts.append(f"{_escape(name)}{_ASSIGN}{_escape(str(partition[name]))}") + return _SEPARATOR.join(parts) diff --git a/tests/stateclient/__init__.py b/tests/stateclient/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/stateclient/keys.py b/tests/stateclient/keys.py new file mode 100644 index 0000000..a29f335 --- /dev/null +++ b/tests/stateclient/keys.py @@ -0,0 +1,33 @@ +from cloudquery.sdk.stateclient import state_key + + +def test_table_key(): + assert state_key("items") == "items" + + +def test_partition_fields_are_sorted(): + assert state_key("t", b="2", a=1) == state_key("t", a=1, b="2") == "t|a=1|b=2" + + +def test_client_id_comes_before_partition_fields(): + assert state_key("t", client_id="acct", a="x") == "t|client_id=acct|a=x" + + +def test_separators_are_escaped(): + assert state_key("t", a="x|b=y") == "t|a=x\\|b\\=y" + assert state_key("t", a="x|b=y") != state_key("t", a="x", b="y") + assert state_key("t", a="x\\") != state_key("t", a="x\\\\") + + +def test_tables_clients_and_partitions_never_share_a_key(): + keys = [ + state_key("t"), + state_key("t|a=1"), + state_key("t", a=1), + state_key("t", client_id="a=1"), + state_key("t", client_id="c", a=1), + state_key("t", client_id="c|a=1"), + state_key("t", a="1|b=2"), + state_key("t", a=1, b=2), + ] + assert len(set(keys)) == len(keys) From 1b4fde9a07ced60ae88b0f4e707ecee83c277f67 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 14:17:34 +0300 Subject: [PATCH 4/6] 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 5/6] 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) From 301f2786c0f1abc1bd506f1a68a35d1aff9b78d7 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 20 Oct 2026 16:05:41 +0300 Subject: [PATCH 6/6] chore: Lock cloudquery-plugin-pb 0.0.56, exempt from exclude-newer --- pyproject.toml | 1 + uv.lock | 13 ++++++++----- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 13d7df4..9da402d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,3 +61,4 @@ python_files = "tests/*.py" [tool.uv] exclude-newer = "1 week" +exclude-newer-package = { cloudquery-plugin-pb = false } diff --git a/uv.lock b/uv.lock index 28c8762..1f75c29 100644 --- a/uv.lock +++ b/uv.lock @@ -11,9 +11,12 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-08T08:49:57.565562033Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P1W" +[options.exclude-newer-package] +cloudquery-plugin-pb = false + [[package]] name = "black" version = "26.3.1" @@ -65,7 +68,7 @@ wheels = [ [[package]] name = "cloudquery-plugin-pb" -version = "0.0.54" +version = "0.0.56" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, @@ -73,9 +76,9 @@ dependencies = [ { name = "protobuf" }, { name = "pyarrow" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/91/b8/15852c6d8e6f4e68780edbf974640a6eb90219f3675c1f86cd82f811fb85/cloudquery_plugin_pb-0.0.54.tar.gz", hash = "sha256:c17ddcdb959fd6f0d85ec78d3d94f21060d32465ba50fcac4261def03d0a4eee", size = 16964, upload-time = "2026-03-13T09:57:54.769Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/33/e3ac909a1024e574a0ebb3e9943d22175411cc3d93a143fb8d838053ffa3/cloudquery_plugin_pb-0.0.56.tar.gz", hash = "sha256:95ea8a9948d3d8fff2d985bc2b427cc129dbc3eaef9446b18a543816649bb367", size = 18615, upload-time = "2026-10-20T08:08:51.925Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/eb/498a0f706f746c6ff9173ae214bd098f6d03f4386547b9668e332325e8b0/cloudquery_plugin_pb-0.0.54-py2.py3-none-any.whl", hash = "sha256:587586ef045f6da70348ca1d2275add2f4e76101d50667f8ebd3a86f6299c4df", size = 19524, upload-time = "2026-03-13T09:57:53.402Z" }, + { url = "https://files.pythonhosted.org/packages/3b/72/e2d245bd38d77840b1da138105d0792037cbe5a2113a0ffb52c9372bac91/cloudquery_plugin_pb-0.0.56-py3-none-any.whl", hash = "sha256:529f582cb7afff09e2aee41c093068bfb0d27251e73f83e3479eea345bd6015f", size = 20877, upload-time = "2026-10-20T08:08:50.723Z" }, ] [[package]] @@ -109,7 +112,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "cloudquery-plugin-pb", specifier = ">=0.0.54" }, + { name = "cloudquery-plugin-pb", specifier = ">=0.0.55" }, { name = "grpcio", specifier = ">=1.78.0" }, { name = "grpcio-tools", specifier = ">=1.78.0" }, { name = "jinja2", specifier = ">=3.1.6" },