fix: Address state checkpoint review findings

cloudquery/sdk/internal/checkpoint.py

@@ -13,10 +13,11 @@
 
 from __future__ import annotations
 
+import contextvars
 import queue
 import threading
 from contextlib import contextmanager
-from typing import Dict, Iterable, Iterator, List, Optional, Set
+from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple
 
 from cloudquery.sdk.message import SyncCheckpointMessage
 
@@ -29,8 +30,10 @@ class Scope:
     spawned is finished.
     """
 
-    def __init__(self, parent: Optional[Scope]):
+    def __init__(self, parent: Optional[Scope], sync: Optional[object]):
         self.parent = parent
+        # the sync the scope belongs to, so that resolvers still running after it ended cannot checkpoint into the next one
+        self.sync = sync
         # child scopes that are not finished yet
         self.open_children: Set[Scope] = set()
         self.resolving = True
@@ -44,12 +47,12 @@ def finished(self) -> bool:
 class _Pending:
     """A checkpoint waiting for scopes to finish and for earlier checkpoints of its keys to be sent."""
 
-    def __init__(self, msg: SyncCheckpointMessage, keys: List[str]):
+    def __init__(self, msg: SyncCheckpointMessage, keys: List[str], sync: object):
         self.msg = msg
         self.keys = keys
+        self.sync = sync
         self.blockers = 0
         self.dependents: List[_Pending] = []
-        self.sent = False
 
 
 class Checkpoints:
@@ -62,44 +65,52 @@ class Checkpoints:
 
     def __init__(self):
         self._lock = threading.Lock()
-        self._local = threading.local()
+        self._scope: contextvars.ContextVar[Optional[Scope]] = contextvars.ContextVar(
+            "cloudquery_checkpoint_scope", default=None
+        )
         self._res: Optional[queue.Queue] = None
         # newest unsent checkpoint of each key
         self._unsent: Dict[str, _Pending] = {}
+        self._sync: Optional[object] = None
 
     def start(self, res: queue.Queue) -> None:
         """Sends later checkpoints on res. Called by the scheduler when a sync starts."""
         with self._lock:
             self._res = res
+            self._sync = object()
             self._unsent = {}
 
     def stop(self) -> None:
         """Drops unsent checkpoints and rejects later ones. Called by the scheduler when a sync ends."""
         with self._lock:
             self._res = None
+            self._sync = None
             self._unsent = {}
 
     def new_scope(self, parent: Optional[Scope]) -> Scope:
         """Returns a scope for a table resolver call, before the call is submitted."""
-        scope = Scope(parent)
+        scope = Scope(parent, parent.sync if parent is not None else self._sync)
         if parent is not None:
             with self._lock:
                 parent.open_children.add(scope)
         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
+        """
+        Makes scope the current one in this context, for the state changes its resolver makes.
+        To change state from other threads, run them in a copy of the context, e.g. with
+        ``contextvars.copy_context().run``; asyncio tasks copy it already.
+        """
+        token = self._scope.set(scope)
         try:
             yield
         finally:
-            self._local.scope = previous
+            self._scope.reset(token)
 
     def current_scope(self) -> Optional[Scope]:
-        """Returns the scope of the table resolver running on this thread, or None outside table resolvers."""
-        return getattr(self._local, "scope", None)
+        """Returns the scope of the table resolver running in this context, or None outside table resolvers."""
+        return self._scope.get()
 
     def resolver_finished(self, scope: Scope) -> None:
         """
@@ -121,48 +132,62 @@ def resolver_finished(self, scope: Scope) -> None:
 
     def schedule(
         self,
-        scopes: Iterable[Optional[Scope]],
-        keys: List[str],
-        msg: SyncCheckpointMessage,
+        checkpoints: List[
+            Tuple[Iterable[Optional[Scope]], List[str], SyncCheckpointMessage]
+        ],
     ) -> None:
         """
-        Sends msg once every resource the scopes produced so far is on the queue, including
-        the relations spawned from them, and after every earlier unsent checkpoint of keys.
-        A None scope stands for changes made outside table resolvers, which wait for nothing.
+        Sends each (scopes, keys, msg) once every resource the scopes produced so far is on the
+        queue, including the relations spawned from them, and after every earlier unsent
+        checkpoint of keys. A None scope stands for changes made outside table resolvers, which
+        wait for nothing. Schedules all of them or, if it raises, none of them.
 
         Raises:
-            RuntimeError: if no sync is running, as the checkpoint could not be sent.
+            RuntimeError: if no sync is running, as the checkpoints could not be sent.
         """
         with self._lock:
             if self._res is None:
                 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])
+            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()
-            pending.sent = True
-            if self._res is not None:
+            if pending.sync is self._sync:
                 self._res.put(pending.msg)
             for key in pending.keys:
                 if self._unsent.get(key) is pending:

cloudquery/sdk/scheduler/scheduler.py

@@ -9,7 +9,6 @@
     SyncInsertMessage,
     SyncMigrateTableMessage,
     SyncErrorMessage,
-    SyncCheckpointMessage,
 )
 from cloudquery.sdk.internal.checkpoint import Checkpoints, Scope
 from cloudquery.sdk.schema import Resource
@@ -35,6 +34,11 @@ def __init__(self) -> None:
         pass
 
 
+class TableResolversSubmitted:
+    def __init__(self) -> None:
+        pass
+
+
 class Scheduler:
     def __init__(
         self, concurrency: int, queue_size: int = 0, max_depth: int = 3, logger=None
@@ -95,8 +99,7 @@ def resolve_table(
         res: queue.Queue,
         scope: Optional[Scope] = None,
     ):
-        if scope is None:
-            scope = self.checkpoints.new_scope(None)
+        scope = scope or self.checkpoints.new_scope(None)
         try:
             with self.checkpoints.entered(scope):
                 self._resolve_table(resolver, depth, client, parent_item, res, scope)
@@ -197,19 +200,16 @@ def _sync(
         res: queue.Queue,
         deterministic_cq_id=False,
     ):
-        for resolver in resolvers:
-            clients = resolver.multiplex(client)
-            for client in clients:
-                res.put(TableResolverStarted())
-                self._pools[0].submit(
-                    self.resolve_table,
-                    resolver,
-                    0,
-                    client,
-                    None,
-                    res,
-                    self.checkpoints.new_scope(None),
-                )
+        try:
+            for resolver in resolvers:
+                clients = resolver.multiplex(client)
+                for client in clients:
+                    res.put(TableResolverStarted())
+                    self._pools[0].submit(
+                        self.resolve_table, resolver, 0, client, None, res
+                    )
+        finally:
+            res.put(TableResolversSubmitted())
 
     def sync(
         self, client, resolvers: List[TableResolver], deterministic_cq_id=False
@@ -221,31 +221,24 @@ def sync(
         try:
             thread = futures.ThreadPoolExecutor()
             thread.submit(self._sync, client, resolvers, res, deterministic_cq_id)
+            submitted = False
             total_table_resolvers = 0
             finished_table_resolvers = 0
-            while True:
+            while not (submitted and total_table_resolvers == finished_table_resolvers):
                 message = res.get()
-                if isinstance(message, TableResolverStarted):
+                if isinstance(message, TableResolversSubmitted):
+                    submitted = True
+                elif isinstance(message, TableResolverStarted):
                     total_table_resolvers += 1
-                    if total_table_resolvers == finished_table_resolvers:
-                        break
-                    continue
-                if isinstance(message, TableResolverFinished):
+                elif isinstance(message, TableResolverFinished):
                     finished_table_resolvers += 1
-                    if total_table_resolvers == finished_table_resolvers:
-                        break
-                    continue
-                yield message
+                else:
+                    yield message
 
             self._post_sync_hook()
             # checkpoints the post-sync hook flushed, after every resource
-            while True:
-                try:
-                    message = res.get_nowait()
-                except queue.Empty:
-                    break
-                if isinstance(message, SyncCheckpointMessage):
-                    yield message
+            while not res.empty():
+                yield res.get_nowait()
             thread.shutdown(wait=True)
         finally:
             self.checkpoints.stop()

cloudquery/sdk/stateclient/keys.py

@@ -18,7 +18,7 @@ def _escape(part: str) -> str:
 def state_key(
     table: str,
     *,
-    client_id: Optional[str] = None,
+    client_id: Optional[Union[str, int]] = None,
     **partition: Union[str, int],
 ) -> str:
     """
@@ -38,7 +38,7 @@ def state_key(
     """
     parts = [_escape(table)]
     if client_id is not None:
-        parts.append(f"client_id{_ASSIGN}{_escape(client_id)}")
+        parts.append(f"client_id{_ASSIGN}{_escape(str(client_id))}")
     for name in sorted(partition):
         parts.append(f"{_escape(name)}{_ASSIGN}{_escape(str(partition[name]))}")
     return _SEPARATOR.join(parts)

cloudquery/sdk/stateclient/stateclient.py

@@ -120,9 +120,18 @@ def _send_checkpoints(
         groups: Dict[FrozenSet[Optional[Scope]], List[str]] = {}
         for key, scopes in changes.items():
             groups.setdefault(frozenset(scopes), []).append(key)
-        for scopes, keys in groups.items():
-            record = self.table.record([(key, values[key]) for key in keys])
-            self.checkpoints.schedule(scopes, keys, SyncCheckpointMessage(record))
+        self.checkpoints.schedule(
+            [
+                (
+                    scopes,
+                    keys,
+                    SyncCheckpointMessage(
+                        self.table.record([(key, values[key]) for key in keys])
+                    ),
+                )
+                for scopes, keys in groups.items()
+            ]
+        )
 
     @abstractmethod
     def migrate_state_table(self):