fix: Address checkpoint review findings

- Wait for _sync to finish submitting resolvers before ending the sync
- Yield every message left after the post-sync hook
- Track the resolver scope in a ContextVar so copied contexts keep it
- Schedule a flush's checkpoints all or none
- Accept non-string client ids in state_key
- Drop unused _Pending.sent and the redundant root scope in _sync

Author: erezrokah

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
 
@@ -49,7 +50,6 @@ def __init__(self, msg: SyncCheckpointMessage, keys: List[str]):
         self.keys = keys
         self.blockers = 0
         self.dependents: List[_Pending] = []
-        self.sent = False
 
 
 class Checkpoints:
@@ -62,7 +62,9 @@ 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] = {}
@@ -89,17 +91,20 @@ def new_scope(self, parent: Optional[Scope]) -> 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,47 +126,56 @@ 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, 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)
+        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:

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):

tests/scheduler/checkpoints.py

@@ -1,6 +1,8 @@
+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
@@ -300,3 +302,53 @@ def test_builder_uses_checkpoints_only_when_the_cli_supports_them(monkeypatch):
     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

tests/stateclient/keys.py

@@ -31,3 +31,7 @@ def test_tables_clients_and_partitions_never_share_a_key():
         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"