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.

Author: jpollard-cs

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 = []

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(

cloudquery/sdk/message/__init__.py

@@ -3,6 +3,7 @@
     SyncInsertMessage,
     SyncMigrateTableMessage,
     SyncErrorMessage,
+    SyncCheckpointMessage,
 )
 from .write import (
     WriteMessage,

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

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