diff --git a/cloudquery/sdk/internal/checkpoint.py b/cloudquery/sdk/internal/checkpoint.py new file mode 100644 index 0000000..394d6a3 --- /dev/null +++ b/cloudquery/sdk/internal/checkpoint.py @@ -0,0 +1,199 @@ +""" +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 contextvars +import queue +import threading +from contextlib import contextmanager +from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple + +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], 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 + # 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], sync: object): + self.msg = msg + self.keys = keys + self.sync = sync + self.blockers = 0 + self.dependents: List[_Pending] = [] + + +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._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] = {} + 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, parent.sync if parent is not None else self._sync) + 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 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._scope.reset(token) + + def current_scope(self) -> Optional[Scope]: + """Returns the scope of the table resolver running in this context, or None outside table resolvers.""" + return self._scope.get() + + def 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, + checkpoints: List[ + Tuple[Iterable[Optional[Scope]], List[str], SyncCheckpointMessage] + ], + ) -> None: + """ + 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 checkpoints 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" + ) + 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) + + def _schedule( + self, + scopes: Iterable[Optional[Scope]], + keys: List[str], + msg: SyncCheckpointMessage, + ) -> None: + 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 + 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() + if pending.sync is self._sync: + 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..6ea9abe 100644 --- a/cloudquery/sdk/scheduler/scheduler.py +++ b/cloudquery/sdk/scheduler/scheduler.py @@ -10,6 +10,7 @@ SyncMigrateTableMessage, SyncErrorMessage, ) +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 @@ -33,12 +34,19 @@ 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 ): 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 +97,27 @@ def resolve_table( client: Client, parent_item: Optional[Resource], res: queue.Queue, + scope: Optional[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) + 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 +133,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 +163,7 @@ def resolve_table( client, resource, res, + self.checkpoints.new_scope(scope), ) total_resources += 1 if depth == 0: @@ -164,8 +192,6 @@ def resolve_table( exc_info=e, ) res.put(SyncErrorMessage(resolver.table.name, str(e))) - finally: - res.put(TableResolverFinished()) def _sync( self, @@ -174,13 +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 - ) + 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 @@ -188,26 +217,31 @@ 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._post_sync_hook() - thread.shutdown(wait=True) + self.checkpoints.start(res) + 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 not (submitted and total_table_resolvers == finished_table_resolvers): + message = res.get() + if isinstance(message, TableResolversSubmitted): + submitted = True + elif isinstance(message, TableResolverStarted): + total_table_resolvers += 1 + elif isinstance(message, TableResolverFinished): + finished_table_resolvers += 1 + else: + yield message + + self._post_sync_hook() + # checkpoints the post-sync hook flushed, after every resource + while not res.empty(): + yield res.get_nowait() + thread.shutdown(wait=True) + finally: + self.checkpoints.stop() def _send_migrate_table_messages( self, resolvers: List[TableResolver] 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..7e14c6c --- /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[Union[str, int]] = 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(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 0b0dded..dfa245a 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 backend_options.supports_checkpoints: + checkpoints = None + return ConnectedStateClient( + backend_options=backend_options, checkpoints=checkpoints + ) class StateClient(ABC): @@ -43,31 +59,79 @@ 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) + 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): @@ -93,6 +157,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 +211,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 +258,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..9da402d 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", @@ -61,3 +61,4 @@ python_files = "tests/*.py" [tool.uv] exclude-newer = "1 week" +exclude-newer-package = { cloudquery-plugin-pb = false } diff --git a/tests/scheduler/checkpoints.py b/tests/scheduler/checkpoints.py new file mode 100644 index 0000000..61fc182 --- /dev/null +++ b/tests/scheduler/checkpoints.py @@ -0,0 +1,397 @@ +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 +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 + + +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 + + +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) 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 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..d431b08 --- /dev/null +++ b/tests/stateclient/keys.py @@ -0,0 +1,37 @@ +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) + + +def test_client_id_that_is_not_a_string(): + assert state_key("t", client_id=123) == "t|client_id=123" 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" },