cloudquery-plugin-sdk 0.1.55__py3-none-any.whl

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/sdk/internal/memdb/__init__.py

@@ -0,0 +1 @@
+from .memdb import MemDB

cloudquery/sdk/internal/memdb/memdb.py

@@ -0,0 +1,221 @@
+import json
+
+from cloudquery.sdk import plugin
+from cloudquery.sdk import message
+from cloudquery.sdk import schema
+from cloudquery.sdk.scheduler import Scheduler, TableResolver
+from typing import List, Generator, Dict, Any
+import pyarrow as pa
+from cloudquery.sdk.schema.table import Table
+from cloudquery.sdk.schema.arrow import METADATA_TABLE_NAME
+from cloudquery.sdk.types import JSONType
+from dataclasses import dataclass, field
+
+NAME = "memdb"
+VERSION = "development"
+
+
+class Client:
+    def __init__(self) -> None:
+        pass
+
+    def id(self):
+        return "memdb"
+
+
+class MemDBResolver(TableResolver):
+    def __init__(
+        self, table: Table, records: List, child_resolvers: list[TableResolver] = None
+    ) -> None:
+        super().__init__(table=table, child_resolvers=child_resolvers)
+        self._records = records
+
+    def resolve(self, client: None, parent_resource) -> Generator[Any, None, None]:
+        for record in self._records:
+            yield record
+
+
+class Table1Relation1(Table):
+    def __init__(self) -> None:
+        super().__init__(
+            name="table_1_relation_1",
+            columns=[
+                schema.Column(
+                    name="name",
+                    type=pa.string(),
+                    primary_key=True,
+                    not_null=True,
+                    unique=True,
+                ),
+                schema.Column(name="data", type=JSONType()),
+            ],
+            title="Table 1 Relation 1",
+            description="Test Table 1 Relation 1",
+        )
+
+    @property
+    def resolver(self):
+        return MemDBResolver(
+            self,
+            records=[
+                {"name": "a", "data": {"a": 1}},
+                {"name": "b", "data": {"b": 2}},
+                {"name": "c", "data": {"c": 3}},
+            ],
+        )
+
+
+class Table1(Table):
+    def __init__(self) -> None:
+        super().__init__(
+            name="table_1",
+            columns=[
+                schema.Column(
+                    name="name",
+                    type=pa.string(),
+                    primary_key=True,
+                    not_null=True,
+                    unique=True,
+                ),
+                schema.Column(
+                    name="id",
+                    type=pa.int64(),
+                    primary_key=True,
+                    not_null=True,
+                    unique=True,
+                    incremental_key=True,
+                ),
+            ],
+            title="Table 1",
+            description="Test Table 1",
+            is_incremental=True,
+            relations=[Table1Relation1()],
+        )
+
+    @property
+    def resolver(self):
+        child_resolvers: list[TableResolver] = []
+        for rel in self.relations:
+            child_resolvers.append(rel.resolver)
+
+        return MemDBResolver(
+            self,
+            records=[
+                {"name": "a", "id": 1},
+                {"name": "b", "id": 2},
+                {"name": "c", "id": 3},
+            ],
+            child_resolvers=child_resolvers,
+        )
+
+
+class Table2(Table):
+    def __init__(self) -> None:
+        super().__init__(
+            name="table_2",
+            columns=[
+                schema.Column(
+                    name="name",
+                    type=pa.string(),
+                    primary_key=True,
+                    not_null=True,
+                    unique=True,
+                ),
+                schema.Column(name="id", type=pa.int64()),
+            ],
+            title="Table 2",
+            description="Test Table 2",
+        )
+
+    @property
+    def resolver(self):
+        return MemDBResolver(
+            self,
+            records=[
+                {"name": "a", "id": 1},
+                {"name": "b", "id": 2},
+                {"name": "c", "id": 3},
+            ],
+        )
+
+
+@dataclass
+class Spec:
+    concurrency: int = field(default=1000)
+    queue_size: int = field(default=1000)
+
+
+class MemDB(plugin.Plugin):
+    def __init__(self) -> None:
+        super().__init__(
+            NAME, VERSION, opts=plugin.plugin.Options(team="cloudquery", kind="source")
+        )
+        table1 = Table1()
+        table2 = Table2()
+        self._tables: Dict[str, schema.Table] = {
+            table1.name: table1,
+            table2.name: table2,
+        }
+        self._db: List[pa.RecordBatch] = []
+        self._client = Client()
+
+    def set_logger(self, logger) -> None:
+        self._logger = logger
+
+    def init(self, spec, no_connection: bool = False):
+        if no_connection:
+            return
+        self._spec_json = json.loads(spec)
+        self._spec = Spec(**self._spec_json)
+        self._scheduler = Scheduler(
+            concurrency=self._spec.concurrency,
+            queue_size=self._spec.queue_size,
+            logger=self._logger,
+        )
+
+    def get_tables(self, options: plugin.TableOptions = None) -> List[plugin.Table]:
+        tables = list(self._tables.values())
+
+        # set parent table relationships
+        for table in tables:
+            for relation in table.relations:
+                relation.parent = table
+
+        return schema.filter_dfs(tables, options.tables, options.skip_tables)
+
+    def sync(
+        self, options: plugin.SyncOptions
+    ) -> Generator[message.SyncMessage, None, None]:
+        resolvers: list[TableResolver] = []
+        for table in self.get_tables(
+            plugin.TableOptions(
+                tables=options.tables,
+                skip_tables=options.skip_tables,
+                skip_dependent_tables=options.skip_dependent_tables,
+            )
+        ):
+            resolvers.append(table.resolver)
+
+        return self._scheduler.sync(
+            self._client, resolvers, options.deterministic_cq_id
+        )
+
+    def write(self, writer: Generator[message.WriteMessage, None, None]) -> None:
+        for msg in writer:
+            if isinstance(msg, message.WriteMigrateTableMessage):
+                pass
+            elif isinstance(msg, message.WriteInsertMessage):
+                self._db.append(msg.record)
+            else:
+                raise NotImplementedError(f"Unknown message type {type(msg)}")
+
+    def read(self, table: Table) -> Generator[message.ReadMessage, None, None]:
+        for record in self._db:
+            recordMetadata = record.schema.metadata.get(METADATA_TABLE_NAME).decode(
+                "utf-8"
+            )
+            if recordMetadata == table.name:
+                yield message.ReadMessage(record)
+
+    def close(self) -> None:
+        self._db = {}

cloudquery/sdk/internal/servers/discovery_v1/discovery.py

@@ -0,0 +1,11 @@
+from typing import List
+
+from cloudquery.discovery_v1 import discovery_pb2, discovery_pb2_grpc
+
+
+class DiscoveryServicer(discovery_pb2_grpc.DiscoveryServicer):
+    def __init__(self, versions: List[int]):
+        self._versions = versions
+
+    def GetVersions(self, request, context):
+        return discovery_pb2.GetVersions.Response(versions=self._versions)

cloudquery/sdk/internal/servers/plugin_v3/__init__.py

@@ -0,0 +1 @@
+from .plugin import PluginServicer

cloudquery/sdk/internal/servers/plugin_v3/plugin.py

@@ -0,0 +1,156 @@
+from typing import Generator
+
+import pyarrow as pa
+import structlog
+from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc, arrow
+
+from cloudquery.sdk.message import (
+    SyncInsertMessage,
+    SyncMigrateTableMessage,
+    SyncErrorMessage,
+    SyncCheckpointMessage,
+    WriteInsertMessage,
+    WriteMigrateTableMessage,
+    WriteMessage,
+    WriteDeleteStale,
+)
+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
+
+
+class PluginServicer(plugin_pb2_grpc.PluginServicer):
+    def __init__(self, plugin: Plugin, logger=None):
+        self._logger = logger if logger is not None else structlog.get_logger()
+        self._plugin = plugin
+
+    def GetName(self, request, context):
+        return plugin_pb2.GetName.Response(name=self._plugin.name())
+
+    def GetVersion(self, request, context):
+        return plugin_pb2.GetVersion.Response(version=self._plugin.version())
+
+    def GetSpecSchema(self, request, context):
+        return plugin_pb2.GetSpecSchema.Response(json_schema=self._plugin.json_schema())
+
+    def Init(self, request: plugin_pb2.Init.Request, context):
+        self._plugin.init(
+            sanitize_spec(request.spec), no_connection=request.no_connection
+        )
+        return plugin_pb2.Init.Response()
+
+    def GetTables(self, request: plugin_pb2.GetTables.Request, context):
+        tables = self._plugin.get_tables(
+            TableOptions(
+                tables=request.tables,
+                skip_tables=request.skip_tables,
+                skip_dependent_tables=request.skip_dependent_tables,
+            )
+        )
+        flattened_tables = flatten_tables(tables)
+        schema = tables_to_arrow_schemas(flattened_tables)
+        tablesBytes = []
+        for s in schema:
+            sink = pa.BufferOutputStream()
+            writer = pa.ipc.new_stream(sink, s)
+            writer.close()
+            buf = sink.getvalue().to_pybytes()
+            tablesBytes.append(buf)
+
+        return plugin_pb2.GetTables.Response(tables=tablesBytes)
+
+    def Sync(self, request, context):
+        options = SyncOptions(
+            deterministic_cq_id=False,  # TODO
+            skip_dependent_tables=request.skip_dependent_tables,
+            skip_tables=request.skip_tables,
+            tables=request.tables,
+            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):
+            if isinstance(msg, SyncInsertMessage):
+                buf = arrow.record_to_bytes(msg.record)
+                yield plugin_pb2.Sync.Response(
+                    insert=plugin_pb2.Sync.MessageInsert(record=buf)
+                )
+            elif isinstance(msg, SyncMigrateTableMessage):
+                buf = arrow.schema_to_bytes(msg.table)
+                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(
+                        table_name=msg.table_name, error=msg.error
+                    )
+                )
+            else:
+                # unknown sync message type
+                raise NotImplementedError()
+
+    def Read(
+        self, request: plugin_pb2.Read.Request, context
+    ) -> Generator[plugin_pb2.Read.Response, None, None]:
+        schema = arrow.new_schema_from_bytes(request.table)
+        table = Table.from_arrow_schema(schema)
+        for msg in self._plugin.read(table):
+            buf = arrow.record_to_bytes(msg.record)
+            yield plugin_pb2.Read.Response(record=buf)
+
+    def Write(
+        self, request_iterator: Generator[plugin_pb2.Write.Request, None, None], context
+    ):
+        def msg_iterator() -> Generator[WriteMessage, None, None]:
+            for msg in request_iterator:
+                field = msg.WhichOneof("message")
+                if field == "migrate_table":
+                    sc = arrow.new_schema_from_bytes(msg.migrate_table.table)
+                    table = Table.from_arrow_schema(sc)
+                    yield WriteMigrateTableMessage(
+                        table=table, migrate_force=msg.migrate_table.migrate_force
+                    )
+                elif field == "insert":
+                    yield WriteInsertMessage(
+                        record=arrow.new_record_from_bytes(msg.insert.record)
+                    )
+                elif field == "delete":
+                    yield WriteDeleteStale(
+                        table_name=msg.delete.table_name,
+                        source_name=msg.delete.source_name,
+                        sync_time=msg.delete.sync_time.ToDatetime(),
+                    )
+                elif field is None:
+                    continue
+                else:
+                    raise NotImplementedError(f"unknown write message type {field}")
+
+        self._plugin.write(msg_iterator())
+        return plugin_pb2.Write.Response()
+
+    def Close(self, request, context):
+        self._plugin.close()
+        return plugin_pb2.Close.Response()
+
+
+def sanitize_spec(spec=None):
+    return b"{}" if spec is None or spec == b"" or spec == b"null" else spec

cloudquery/sdk/message/__init__.py

@@ -0,0 +1,14 @@
+from .sync import (
+    SyncMessage,
+    SyncInsertMessage,
+    SyncMigrateTableMessage,
+    SyncErrorMessage,
+    SyncCheckpointMessage,
+)
+from .write import (
+    WriteMessage,
+    WriteInsertMessage,
+    WriteMigrateTableMessage,
+    WriteDeleteStale,
+)
+from .read import ReadMessage

cloudquery/sdk/message/read.py

@@ -0,0 +1,6 @@
+import pyarrow as pa
+
+
+class ReadMessage:
+    def __init__(self, record: pa.RecordBatch):
+        self.record = record

cloudquery/sdk/message/sync.py

@@ -0,0 +1,31 @@
+import pyarrow as pa
+
+
+class SyncMessage:
+    pass
+
+
+class SyncInsertMessage(SyncMessage):
+    def __init__(self, record: pa.RecordBatch):
+        self.record = record
+
+
+class SyncMigrateTableMessage(SyncMessage):
+    def __init__(self, table: pa.Schema):
+        self.table = table
+
+
+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/sdk/message/write.py

@@ -0,0 +1,24 @@
+import pyarrow as pa
+from cloudquery.sdk.schema import Table
+
+
+class WriteMessage:
+    pass
+
+
+class WriteInsertMessage(WriteMessage):
+    def __init__(self, record: pa.RecordBatch):
+        self.record = record
+
+
+class WriteMigrateTableMessage(WriteMessage):
+    def __init__(self, table: Table, migrate_force: bool):
+        self.table = table
+        self.migrate_force = migrate_force
+
+
+class WriteDeleteStale(WriteMessage):
+    def __init__(self, table_name: str, source_name: str, sync_time):
+        self.table_name = table_name
+        self.source_name = source_name
+        self.sync_time = sync_time

cloudquery/sdk/plugin/__init__.py

@@ -0,0 +1 @@
+from .plugin import Plugin, Table, TableOptions, SyncOptions