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

cloudquery_plugin_sdk-0.1.55/PKG-INFO

@@ -0,0 +1,50 @@
+Metadata-Version: 2.4
+Name: cloudquery-plugin-sdk
+Version: 0.1.55
+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.55}/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.55/cloudquery/sdk/internal/checkpoint.py

@@ -0,0 +1,199 @@
+"""
+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 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] = []
+
+    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()
+
+    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:
+            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"
+                )
+            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()
+            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.55}/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.55}/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.55}/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.55}/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.55}/cloudquery/sdk/scheduler/scheduler.py

@@ -10,6 +10,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
@@ -33,12 +34,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()
@@ -89,7 +97,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,7 +133,6 @@ class Scheduler:
                     table=resolver.table.name,
                     depth=depth,
                 )
-            total_resources = 0
             for item in resolver.resolve(client, parent_item):
                 try:
                     resource = self.resolve_resource(
@@ -136,6 +163,7 @@ class Scheduler:
                         client,
                         resource,
                         res,
+                        self.checkpoints.new_scope(scope),
                     )
                 total_resources += 1
             if depth == 0:
@@ -164,8 +192,6 @@ class Scheduler:
                 exc_info=e,
             )
             res.put(SyncErrorMessage(resolver.table.name, str(e)))
-        finally:
-            res.put(TableResolverFinished())
 
     def _sync(
         self,
@@ -174,13 +200,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 +217,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.55/cloudquery/sdk/stateclient/__init__.py

@@ -0,0 +1 @@
+from .keys import state_key

cloudquery_plugin_sdk-0.1.55/cloudquery/sdk/stateclient/keys.py

@@ -0,0 +1,44 @@
+"""Builds state keys, so that plugins choose how finely they track progress without inventing key formats."""
+
+from typing import Optional, Union
+
+_SEPARATOR = "|"
+_ASSIGN = "="
+_ESCAPE = "\\"
+
+
+def _escape(part: str) -> str:
+    return (
+        part.replace(_ESCAPE, _ESCAPE * 2)
+        .replace(_SEPARATOR, _ESCAPE + _SEPARATOR)
+        .replace(_ASSIGN, _ESCAPE + _ASSIGN)
+    )
+
+
+def state_key(
+    table: str,
+    *,
+    client_id: Optional[Union[str, int]] = None,
+    **partition: Union[str, int],
+) -> str:
+    """
+    Returns the state key of a table's progress, optionally for one client and one partition of the table.
+
+    A table with one cursor uses ``state_key("items")``. A table that tracks progress per parent, account or
+    other partition passes the fields that identify it, e.g. ``state_key("responses", form_id=form["id"])``.
+    Each such key is its own state row: with checkpoints, a partition's state waits only for the resources of
+    the table resolver that set it, and the CLI persists the keys changed since its last commit together.
+
+    Keys are readable and deterministic: partition fields are sorted, and separators in names and values are
+    escaped, so different tables, clients or partitions never share a key.
+
+    Example:
+        >>> state_key("responses", client_id="acct-1", form_id=123)
+        'responses|client_id=acct-1|form_id=123'
+    """
+    parts = [_escape(table)]
+    if client_id is not None:
+        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)