feat: Send state as checkpoints after the resources it covers

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