cloudquery-plugin-sdk 0.1.54__tar.gz → 0.1.56__tar.gz

cloudquery_plugin_sdk-0.1.56/PKG-INFO

@@ -0,0 +1,50 @@
+Metadata-Version: 2.4
+Name: cloudquery-plugin-sdk
+Version: 0.1.56
+Summary: CloudQuery Plugin SDK for Python
+Author-email: CloudQuery LTD <pypi-packages@cloudquery.io>
+License-Expression: MPL-2.0
+Project-URL: Homepage, https://github.com/cloudquery/plugin-sdk-python
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Internet
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: cloudquery-plugin-pb>=0.0.55
+Requires-Dist: grpcio>=1.78.0
+Requires-Dist: grpcio-tools>=1.78.0
+Requires-Dist: Jinja2>=3.1.6
+Requires-Dist: MarkupSafe>=3.0.3
+Requires-Dist: numpy>=2.4.2
+Requires-Dist: packaging>=26.0
+Requires-Dist: pandas>=3.0.0
+Requires-Dist: protobuf>=6.31.1
+Requires-Dist: pyarrow>=23.0.0
+Requires-Dist: python-dateutil>=2.8.1
+Requires-Dist: pytz>=2025.2
+Requires-Dist: six>=1.17.0
+Requires-Dist: structlog>=25.5.0
+Requires-Dist: tomli>=2.4.0
+Requires-Dist: tzdata>=2025.3
+Dynamic: license-file
+
+# CloudQuery Plugin SDK for Python
+
+This is the high-level package to use for developing CloudQuery plugins in Python.
+
+## Installation
+
+```commandline
+uv add cloudquery-plugin-sdk
+```
+
+## Links
+
+- [Developer Guide](https://www.cloudquery.io/docs/developers/creating-new-plugin/python-source)
+- [PyPI](https://pypi.org/project/cloudquery-plugin-sdk/)

{cloudquery_plugin_sdk-0.1.54 → cloudquery_plugin_sdk-0.1.56}/README.md

@@ -5,7 +5,7 @@ This is the high-level package to use for developing CloudQuery plugins in Pytho
 ## Installation
 
 ```commandline
-pip install cloudquery-plugin-sdk
+uv add cloudquery-plugin-sdk
 ```
 
 ## Links

cloudquery_plugin_sdk-0.1.56/cloudquery/sdk/internal/checkpoint.py

@@ -0,0 +1,252 @@
+"""
+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 contextvars
+import queue
+import threading
+from contextlib import contextmanager
+from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple
+
+from cloudquery.sdk.message import SyncCheckpointMessage
+
+
+class SyncNotRunningError(RuntimeError):
+    """
+    Raised when state checkpoints are flushed while their sync is not running, so they cannot be
+    sent: before or after Scheduler.sync, or from a resolver of an ended sync. Flush state that
+    summarizes the whole sync in the scheduler's post-sync hook instead.
+    """
+
+
+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], 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
+        # 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
+
+
+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], sync: object):
+        self.msg = msg
+        self.keys = keys
+        self.sync = sync
+        self.blockers = 0
+        self.dependents: List[_Pending] = []
+
+
+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._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, 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 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._scope.reset(token)
+
+    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 superseded(self, scope: Optional[Scope]) -> bool:
+        """Returns True if scope belongs to a sync that ended while another sync runs now."""
+        return (
+            scope is not None
+            and self._sync is not None
+            and scope.sync is not self._sync
+        )
+
+    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,
+        checkpoints: List[
+            Tuple[Iterable[Optional[Scope]], List[str], SyncCheckpointMessage]
+        ],
+    ) -> None:
+        """
+        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:
+            SyncNotRunningError: if the sync of the checkpoints is not running.
+        """
+        with self._lock:
+            if self._res is None:
+                raise SyncNotRunningError(
+                    "state checkpoints can only be flushed while the scheduler syncs, "
+                    "e.g. from a table resolver or the post-sync hook"
+                )
+            for scopes, _, _ in checkpoints:
+                if any(self.superseded(s) for s in scopes):
+                    raise SyncNotRunningError(
+                        "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}:
+            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)
+            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])
+
+    @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:
+            pending = ready.pop()
+            if pending.sync is self._sync:
+                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_plugin_sdk-0.1.54 → cloudquery_plugin_sdk-0.1.56}/cloudquery/sdk/internal/servers/plugin_v3/plugin.py

@@ -8,12 +8,18 @@ from cloudquery.sdk.message import (
     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 @@ class PluginServicer(plugin_pb2_grpc.PluginServicer):
             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 @@ class PluginServicer(plugin_pb2_grpc.PluginServicer):
                 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_plugin_sdk-0.1.54 → cloudquery_plugin_sdk-0.1.56}/cloudquery/sdk/message/__init__.py

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

{cloudquery_plugin_sdk-0.1.54 → cloudquery_plugin_sdk-0.1.56}/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_plugin_sdk-0.1.54 → cloudquery_plugin_sdk-0.1.56}/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

{cloudquery_plugin_sdk-0.1.54 → cloudquery_plugin_sdk-0.1.56}/cloudquery/sdk/scheduler/scheduler.py

@@ -1,4 +1,5 @@
 import queue
+import threading
 from concurrent import futures
 from typing import List, Generator, Any, Optional
 
@@ -10,6 +11,7 @@ from cloudquery.sdk.message import (
     SyncMigrateTableMessage,
     SyncErrorMessage,
 )
+from cloudquery.sdk.internal.checkpoint import Checkpoints, Scope
 from cloudquery.sdk.schema import Resource
 from cloudquery.sdk.stateclient.stateclient import StateClient
 from .table_resolver import TableResolver, Client
@@ -18,9 +20,22 @@ QUEUE_PER_WORKER = 100
 
 
 class ThreadPoolExecutorWithQueueSizeLimit(futures.ThreadPoolExecutor):
-    def __init__(self, maxsize, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._work_queue = queue.Queue(maxsize=maxsize)
+    def __init__(self, maxsize: int, max_workers: int, *args, **kwargs):
+        super().__init__(max_workers, *args, **kwargs)
+        self.maxsize = maxsize
+        # submit() puts on the work queue while holding a lock that every pool shares, so a
+        # full bounded queue would block the other pools' workers too: limit it before submit
+        self._slots = threading.BoundedSemaphore(maxsize + max_workers)
+
+    def submit(self, fn, /, *args, **kwargs):
+        self._slots.acquire()
+        try:
+            future = super().submit(fn, *args, **kwargs)
+        except BaseException:
+            self._slots.release()
+            raise
+        future.add_done_callback(lambda _: self._slots.release())
+        return future
 
 
 class TableResolverStarted:
@@ -33,12 +48,19 @@ class TableResolverFinished:
         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
     ):
         self._post_sync_hook = lambda: None
         self._queue = queue.Queue()
+        # pass to StateClientBuilder.build so that state flushes are sent as checkpoints, when the CLI supports them
+        self.checkpoints = Checkpoints()
         self._max_depth = max_depth
         if logger is None:
             self._logger = structlog.get_logger()
@@ -53,7 +75,7 @@ class Scheduler:
         )
         self._pools: List[ThreadPoolExecutorWithQueueSizeLimit] = []
         current_depth_concurrency = concurrency
-        current_depth_queue_size = queue_size
+        current_depth_queue_size = self._queue_size
         for _ in range(max_depth + 1):
             self._pools.append(
                 ThreadPoolExecutorWithQueueSizeLimit(
@@ -89,7 +111,27 @@ class Scheduler:
         client: Client,
         parent_item: Optional[Resource],
         res: queue.Queue,
+        scope: Optional[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)
+        finally:
+            # after every resource of the resolver is on the queue, so its checkpoints follow them
+            self.checkpoints.resolver_finished(scope)
+            res.put(TableResolverFinished())
+
+    def _resolve_table(
+        self,
+        resolver: TableResolver,
+        depth: int,
+        client: Client,
+        parent_item: Optional[Resource],
+        res: queue.Queue,
+        scope: Scope,
     ):
+        total_resources = 0
         try:
             if depth == 0:
                 self._logger.info(
@@ -105,38 +147,39 @@ class Scheduler:
                     table=resolver.table.name,
                     depth=depth,
                 )
-            total_resources = 0
             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,
-                    )
+                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(
@@ -164,8 +207,6 @@ class Scheduler:
                 exc_info=e,
             )
             res.put(SyncErrorMessage(resolver.table.name, str(e)))
-        finally:
-            res.put(TableResolverFinished())
 
     def _sync(
         self,
@@ -174,13 +215,16 @@ class Scheduler:
         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
-                )
+        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
@@ -188,26 +232,31 @@ class Scheduler:
         res = queue.Queue()
         yield from self._send_migrate_table_messages(resolvers)
 
-        thread = futures.ThreadPoolExecutor()
-        thread.submit(self._sync, client, resolvers, res, deterministic_cq_id)
-        total_table_resolvers = 0
-        finished_table_resolvers = 0
-        while True:
-            message = res.get()
-            if isinstance(message, TableResolverStarted):
-                total_table_resolvers += 1
-                if total_table_resolvers == finished_table_resolvers:
-                    break
-                continue
-            if isinstance(message, TableResolverFinished):
-                finished_table_resolvers += 1
-                if total_table_resolvers == finished_table_resolvers:
-                    break
-                continue
-            yield message
-
-        self._post_sync_hook()
-        thread.shutdown(wait=True)
+        self.checkpoints.start(res)
+        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 not (submitted and total_table_resolvers == finished_table_resolvers):
+                message = res.get()
+                if isinstance(message, TableResolversSubmitted):
+                    submitted = True
+                elif isinstance(message, TableResolverStarted):
+                    total_table_resolvers += 1
+                elif isinstance(message, TableResolverFinished):
+                    finished_table_resolvers += 1
+                else:
+                    yield message
+
+            self._post_sync_hook()
+            # checkpoints the post-sync hook flushed, after every resource
+            while not res.empty():
+                yield res.get_nowait()
+            thread.shutdown(wait=True)
+        finally:
+            self.checkpoints.stop()
 
     def _send_migrate_table_messages(
         self, resolvers: List[TableResolver]

cloudquery_plugin_sdk-0.1.56/cloudquery/sdk/stateclient/__init__.py

@@ -0,0 +1,2 @@
+from cloudquery.sdk.internal.checkpoint import SyncNotRunningError
+from .keys import state_key