fix: Send checkpoints flushed while a resource is resolved after it and its relations

Author: disq

cloudquery/sdk/internal/checkpoint.py

@@ -39,6 +39,10 @@ def __init__(self, parent: Optional[Scope], sync: Optional[object]):
         self.resolving = True
         # checkpoints waiting for this scope to finish
         self.waiters: List[_Pending] = []
+        # set while a resource is resolved, until it is on the queue and its relations are submitted
+        self.resource_in_flight = False
+        # checkpoints waiting for that resource and its relations
+        self.after_resource: List[_Pending] = []
 
     def finished(self) -> bool:
         return not self.resolving and not self.open_children
@@ -112,6 +116,29 @@ 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()
 
+    @contextmanager
+    def resource(self, scope: Scope) -> Iterator[None]:
+        """
+        Marks a resource of scope as in flight until it is on the queue and its relations are
+        submitted, so that checkpoints flushed meanwhile, e.g. from its resource hooks, follow
+        the resource and its relations.
+        """
+        with self._lock:
+            scope.resource_in_flight = True
+        try:
+            yield
+        finally:
+            with self._lock:
+                scope.resource_in_flight = False
+                ready: List[_Pending] = []
+                for pending in scope.after_resource:
+                    pending.blockers -= 1
+                    self._wait_for_children(pending, scope)
+                    if pending.blockers == 0:
+                        ready.append(pending)
+                scope.after_resource = []
+                self._send(ready)
+
     def resolver_finished(self, scope: Scope) -> None:
         """
         Marks the resolver of scope as returned, once every resource it produced is on the queue.
@@ -167,11 +194,11 @@ def _schedule(
     ) -> 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)
+            if scope.resource_in_flight:
+                scope.after_resource.append(pending)
                 pending.blockers += 1
+            else:
+                self._wait_for_children(pending, scope)
         earlier = set()
         for key in keys:
             previous = self._unsent.get(key)
@@ -183,6 +210,16 @@ def _schedule(
         if pending.blockers == 0:
             self._send([pending])
 
+    @staticmethod
+    def _wait_for_children(pending: _Pending, scope: Scope) -> None:
+        """
+        Makes pending wait for the relations of scope's resources. Its resources are on the queue
+        already, as none is in flight. Called with the lock held.
+        """
+        for child in scope.open_children:
+            child.waiters.append(pending)
+            pending.blockers += 1
+
     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:

cloudquery/sdk/scheduler/scheduler.py

@@ -134,37 +134,38 @@ def _resolve_table(
                     depth=depth,
                 )
             for item in resolver.resolve(client, parent_item):
-                try:
-                    resource = self.resolve_resource(
-                        resolver, client, parent_item, item
-                    )
-                except Exception as e:
-                    self._logger.error(
-                        "failed to resolve resource",
-                        client_id=client.id(),
-                        table=resolver.table.name,
-                        depth=depth,
-                        exc_info=e,
-                    )
-                    self._logger.debug(
-                        "details about resource that failed to resolve",
-                        client_id=client.id(),
-                        table=resolver.table.name,
-                        resource=item,
-                    )
-                    continue
-                res.put(SyncInsertMessage(resource.to_arrow_record()))
-                for child_resolvers in resolver.child_resolvers:
-                    res.put(TableResolverStarted())
-                    self._pools[depth + 1].submit(
-                        self.resolve_table,
-                        child_resolvers,
-                        depth + 1,
-                        client,
-                        resource,
-                        res,
-                        self.checkpoints.new_scope(scope),
-                    )
+                with self.checkpoints.resource(scope):
+                    try:
+                        resource = self.resolve_resource(
+                            resolver, client, parent_item, item
+                        )
+                    except Exception as e:
+                        self._logger.error(
+                            "failed to resolve resource",
+                            client_id=client.id(),
+                            table=resolver.table.name,
+                            depth=depth,
+                            exc_info=e,
+                        )
+                        self._logger.debug(
+                            "details about resource that failed to resolve",
+                            client_id=client.id(),
+                            table=resolver.table.name,
+                            resource=item,
+                        )
+                        continue
+                    res.put(SyncInsertMessage(resource.to_arrow_record()))
+                    for child_resolvers in resolver.child_resolvers:
+                        res.put(TableResolverStarted())
+                        self._pools[depth + 1].submit(
+                            self.resolve_table,
+                            child_resolvers,
+                            depth + 1,
+                            client,
+                            resource,
+                            res,
+                            self.checkpoints.new_scope(scope),
+                        )
                 total_resources += 1
             if depth == 0:
                 self._logger.info(

tests/scheduler/checkpoints.py

@@ -395,3 +395,164 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one():
     ]
     assert cursors
     assert all(cursor.startswith("next-") for cursor in cursors)
+
+
+class HookCursorResolver(TableResolver):
+    """Checkpoints its cursor from a resource hook, before the resource is on the queue."""
+
+    def __init__(self, state: StateClient, hook: str, parents: int):
+        super().__init__(
+            Table("parent", [Column("id", pa.int64())]),
+            child_resolvers=[ChildResolver(2)],
+        )
+        self._state = state
+        self._hook = hook
+        self._parents = parents
+
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+        for parent_id in range(self._parents):
+            yield {"id": parent_id}
+
+    def pre_resource_resolve(self, client, resource):
+        if self._hook == "pre":
+            self._checkpoint(resource)
+
+    def post_resource_resolve(self, client, resource):
+        if self._hook == "post":
+            self._checkpoint(resource)
+
+    def _checkpoint(self, resource):
+        self._state.set_key("hook-cursor", str(resource.item["id"]))
+        self._state.flush()
+
+
+@pytest.mark.parametrize("hook", ["pre", "post"])
+def test_checkpoints_flushed_by_resource_hooks_follow_the_resource_and_its_relations(
+    hook,
+):
+    parents = 10
+    scheduler = Scheduler(4)
+    state = MemoryStateClient(checkpoints=scheduler.checkpoints)
+    try:
+        messages = list(
+            scheduler.sync(TestClient(), [HookCursorResolver(state, hook, parents)])
+        )
+    finally:
+        scheduler.shutdown()
+
+    seen = set()
+    cursors = []
+    for msg in messages:
+        if isinstance(msg, SyncInsertMessage):
+            table = msg.record.schema.metadata[METADATA_TABLE_NAME]
+            row = msg.record.to_pylist()[0]
+            if table == b"parent":
+                seen.add((row["id"],))
+            else:
+                seen.add((row["parent_id"], row["child_id"]))
+        elif isinstance(msg, SyncCheckpointMessage):
+            cursor = int(checkpoint_rows(msg)["hook-cursor"])
+            cursors.append(cursor)
+            for parent_id in range(cursor + 1):
+                assert (parent_id,) in seen, f"cursor {cursor} sent before its parent"
+                for child_id in range(2):
+                    assert (parent_id, child_id) in seen, (
+                        f"cursor {cursor} sent before child {child_id} of parent "
+                        f"{parent_id}"
+                    )
+    assert cursors and cursors[-1] == parents - 1
+
+
+class GrandchildResolver(TableResolver):
+    """Yields slowly, so it still runs after the child resolver that spawned it returned."""
+
+    def __init__(self):
+        super().__init__(
+            Table(
+                "grandchild",
+                [
+                    Column("parent_id", pa.int64()),
+                    Column("child_id", pa.int64()),
+                    Column("grandchild_id", pa.int64()),
+                ],
+            )
+        )
+
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+        for grandchild_id in range(3):
+            time.sleep(random.uniform(0.005, 0.02))
+            yield {**parent_resource.item, "grandchild_id": grandchild_id}
+
+
+class FastChildResolver(TableResolver):
+    def __init__(self):
+        super().__init__(
+            Table(
+                "fast_child",
+                [Column("parent_id", pa.int64()), Column("child_id", pa.int64())],
+            ),
+            child_resolvers=[GrandchildResolver()],
+        )
+
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+        for child_id in range(2):
+            yield {"parent_id": parent_resource.item["id"], "child_id": child_id}
+
+
+class GrandparentResolver(TableResolver):
+    def __init__(self, state: StateClient, parents: int, flush_each: bool):
+        super().__init__(
+            Table("grandparent", [Column("id", pa.int64())]),
+            child_resolvers=[FastChildResolver()],
+        )
+        self._state = state
+        self._parents = parents
+        self._flush_each = flush_each
+
+    def resolve(self, client, parent_resource) -> Generator[Any, None, None]:
+        for parent_id in range(self._parents):
+            yield {"id": parent_id}
+            if self._flush_each:
+                self._state.set_key("grandparent-cursor", str(parent_id))
+                self._state.flush()
+        if not self._flush_each:
+            # the child resolvers return meanwhile, while their grandchildren still run
+            time.sleep(0.01)
+            self._state.set_key("grandparent-cursor", str(self._parents - 1))
+            self._state.flush()
+
+
+@pytest.mark.parametrize("flush_each", [True, False])
+def test_parent_checkpoints_follow_the_relations_of_their_relations(flush_each):
+    parents = 10
+    scheduler = Scheduler(8)
+    state = MemoryStateClient(checkpoints=scheduler.checkpoints)
+    try:
+        messages = list(
+            scheduler.sync(
+                TestClient(), [GrandparentResolver(state, parents, flush_each)]
+            )
+        )
+    finally:
+        scheduler.shutdown()
+
+    seen = set()
+    cursors = []
+    for msg in messages:
+        if (
+            isinstance(msg, SyncInsertMessage)
+            and msg.record.schema.metadata[METADATA_TABLE_NAME] == b"grandchild"
+        ):
+            row = msg.record.to_pylist()[0]
+            seen.add((row["parent_id"], row["child_id"], row["grandchild_id"]))
+        elif isinstance(msg, SyncCheckpointMessage):
+            cursor = int(checkpoint_rows(msg)["grandparent-cursor"])
+            cursors.append(cursor)
+            for parent_id in range(cursor + 1):
+                for child_id in range(2):
+                    for grandchild_id in range(3):
+                        assert (parent_id, child_id, grandchild_id) in seen, (
+                            f"cursor {cursor} sent before grandchild "
+                            f"{(parent_id, child_id, grandchild_id)}"
+                        )
+    assert cursors and cursors[-1] == parents - 1