From 0dc34a5a7976c2b775d02d687bdbe88d1326bbf3 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 19 Oct 2026 14:04:02 -0400 Subject: [PATCH 1/3] 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/3] 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/3] 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)