fix: Add write support for roundtrip test

Author: yevgenypats

cloudquery/sdk/internal/memdb/memdb.py

@@ -11,18 +11,30 @@
 class MemDB(plugin.Plugin):
     def __init__(self) -> None:
         super().__init__(NAME, VERSION)
-        self._tables: List[schema.Table] = [
-            schema.Table("test_table", [schema.Column("test_column", pa.int64())])
-        ]
-        self._memory_db: Dict[str, pa.record] = {
-            "test_table": pa.record_batch([pa.array([1, 2, 3])], names=["test_column"])
-        }
+        self._db: Dict[str, pa.RecordBatch] = {}
+        self._tables: Dict[str, schema.Table] = {}
 
     def get_tables(self, options: plugin.TableOptions = None) -> List[plugin.Table]:
-        return self._tables
+        tables = list(self._tables.values())
+        return schema.filter_dfs(tables, options.tables, options.skip_tables)
 
     def sync(
         self, options: plugin.SyncOptions
     ) -> Generator[message.SyncMessage, None, None]:
-        for table, record in self._memory_db.items():
+        for table, record in self._db.items():
             yield message.SyncInsertMessage(record)
+
+    def write(self, msg_iterator: Generator[message.WriteMessage, None, None]) -> None:
+        for msg in msg_iterator:
+            if type(msg) == message.WriteMigrateTableMessage:
+                if msg.table.name not in self._db:
+                    self._db[msg.table.name] = msg.table
+                    self._tables[msg.table.name] = msg.table
+            elif type(msg) == message.WriteInsertMessage:
+                table = schema.Table.from_arrow_schema(msg.record.schema)
+                self._db[table.name] = msg.record
+            else:
+                raise NotImplementedError(f"Unknown message type {type(msg)}")
+
+    def close(self) -> None:
+        self._db = {}

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

@@ -1,10 +1,18 @@
 import pyarrow as pa
 import structlog
 
+from typing import Generator
 from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc, arrow
-from cloudquery.sdk.message import SyncInsertMessage, SyncMigrateTableMessage
+from cloudquery.sdk.message import (
+    SyncInsertMessage,
+    SyncMigrateTableMessage,
+    WriteInsertMessage,
+    WriteMigrateTableMessage,
+    WriteMessage,
+    WriteDeleteStale,
+)
 from cloudquery.sdk.plugin.plugin import Plugin, SyncOptions, TableOptions
-from cloudquery.sdk.schema import tables_to_arrow_schemas
+from cloudquery.sdk.schema import tables_to_arrow_schemas, Table
 
 
 class PluginServicer(plugin_pb2_grpc.PluginServicer):
@@ -64,8 +72,34 @@ def Sync(self, request, context):
     def Read(self, request, context):
         raise NotImplementedError()
 
-    def Write(self, request_iterator, context):
-        raise NotImplementedError()
+    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)
+                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()

cloudquery/sdk/message/__init__.py

@@ -1 +1,7 @@
 from .sync import SyncMessage, SyncInsertMessage, SyncMigrateTableMessage
+from .write import (
+    WriteMessage,
+    WriteInsertMessage,
+    WriteMigrateTableMessage,
+    WriteDeleteStale,
+)

cloudquery/sdk/message/write.py

@@ -1,15 +1,23 @@
 import pyarrow as pa
+from cloudquery.sdk.schema import Table
 
 
 class WriteMessage:
     pass
 
 
-class InsertMessage(WriteMessage):
+class WriteInsertMessage(WriteMessage):
     def __init__(self, record: pa.RecordBatch):
         self.record = record
 
 
-class MigrateMessage(WriteMessage):
-    def __init__(self, table: pa.Schema):
+class WriteMigrateTableMessage(WriteMessage):
+    def __init__(self, table: Table):
         self.table = table
+
+
+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/plugin.py

@@ -50,5 +50,8 @@ def get_tables(self, options: TableOptions) -> List[Table]:
     def sync(self, options: SyncOptions) -> Generator[message.SyncMessage, None, None]:
         raise NotImplementedError()
 
+    def write(self, writer: Generator[message.WriteMessage, None, None]) -> None:
+        raise NotImplementedError()
+
     def close(self) -> None:
-        pass
+        raise NotImplementedError()

cloudquery/sdk/scheduler/scheduler.py

@@ -103,7 +103,18 @@ def resolve_table(
                 )
             total_resources = 0
             for item in resolver.resolve(client, parent_item):
-                resource = self.resolve_resource(resolver, client, parent_item, item)
+                try:
+                    resource = self.resolve_resource(
+                        resolver, client, parent_item, item
+                    )
+                except Exception as e:
+                    self._logger.error(
+                        "failed to resolve resource",
+                        table=resolver.table.name,
+                        depth=depth,
+                        exception=e,
+                    )
+                    continue
                 res.put(SyncInsertMessage(resource.to_arrow_record()))
                 for child_resolvers in resolver.child_resolvers:
                     self._pools[depth + 1].submit(

cloudquery/sdk/schema/__init__.py

@@ -1,5 +1,5 @@
 from .column import Column
-from .table import Table, tables_to_arrow_schemas
+from .table import Table, tables_to_arrow_schemas, filter_dfs
 from .resource import Resource
 
 # from .table_resolver import TableReso

cloudquery/sdk/schema/arrow.py

@@ -1,9 +1,9 @@
-METADATA_UNIQUE = "cq:extension:unique"
-METADATA_PRIMARY_KEY = "cq:extension:primary_key"
-METADATA_CONSTRAINT_NAME = "cq:extension:constraint_name"
-METADATA_INCREMENTAL = "cq:extension:incremental"
+METADATA_UNIQUE = b"cq:extension:unique"
+METADATA_PRIMARY_KEY = b"cq:extension:primary_key"
+METADATA_CONSTRAINT_NAME = b"cq:extension:constraint_name"
+METADATA_INCREMENTAL = b"cq:extension:incremental"
 
-METADATA_TRUE = "true"
-METADATA_FALSE = "false"
-METADATA_TABLE_NAME = "cq:table_name"
-METADATA_TABLE_DESCRIPTION = "cq:table_description"
+METADATA_TRUE = b"true"
+METADATA_FALSE = b"false"
+METADATA_TABLE_NAME = b"cq:table_name"
+METADATA_TABLE_DESCRIPTION = b"cq:table_description"

cloudquery/sdk/schema/table.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from typing import List, Generator, Any
-
+import fnmatch
 import pyarrow as pa
 
 from cloudquery.sdk.schema import arrow
@@ -53,6 +53,19 @@ def primary_keys(self):
     def incremental_keys(self):
         return [column.name for column in self.columns if column.incremental_key]
 
+    @classmethod
+    def from_arrow_schema(cls, schema: pa.Schema) -> Table:
+        columns = []
+        for field in schema:
+            columns.append(Column.from_arrow_field(field))
+        return cls(
+            name=schema.metadata[arrow.METADATA_TABLE_NAME].decode("utf-8"),
+            columns=columns,
+            description=schema.metadata.get(arrow.METADATA_TABLE_DESCRIPTION).decode(
+                "utf-8"
+            ),
+        )
+
     def to_arrow_schema(self):
         fields = []
         md = {
@@ -74,3 +87,20 @@ def tables_to_arrow_schemas(tables: List[Table]):
     for table in tables:
         schemas.append(table.to_arrow_schema())
     return schemas
+
+
+def filter_dfs(
+    tables: List[Table], include_tables: List[str], skip_tables: List[str]
+) -> List[Table]:
+    filtered: List[Table] = []
+    for table in tables:
+        matched = False
+        for include_table in include_tables:
+            if fnmatch.fnmatch(table.name, include_table):
+                matched = True
+        for skip_table in skip_tables:
+            if fnmatch.fnmatch(table.name, skip_table):
+                matched = False
+        if matched:
+            filtered.append(table)
+    return filtered

tests/internal/memdb/memdb.py

@@ -6,6 +6,6 @@ def test_memdb():
     p = memdb.MemDB()
     p.init(None)
     msgs = []
-    for msg in p.sync(SyncOptions()):
+    for msg in p.sync(SyncOptions(tables=["*"])):
         msgs.append(msg)
-    assert len(msgs) == 1
+    assert len(msgs) == 0

tests/serve/plugin.py

@@ -1,12 +1,22 @@
 import random
 import grpc
 import time
+import pyarrow as pa
 from concurrent import futures
+from cloudquery.sdk.schema import Table, Column
 from cloudquery.sdk import serve
 from cloudquery.sdk import message
 from cloudquery.plugin_v3 import plugin_pb2_grpc, plugin_pb2, arrow
 from cloudquery.sdk.internal.memdb import MemDB
 
+test_table = Table(
+    "test",
+    [
+        Column("id", pa.int64()),
+        Column("name", pa.string()),
+    ],
+)
+
 
 def test_plugin_serve():
     p = MemDB()
@@ -27,11 +37,31 @@ def test_plugin_serve():
             response = stub.Init(plugin_pb2.Init.Request(spec=b""))
             assert response is not None
 
-            response = stub.GetTables(plugin_pb2.GetTables.Request())
+            def writer_iterator():
+                buf = arrow.schema_to_bytes(test_table.to_arrow_schema())
+                yield plugin_pb2.Write.Request(
+                    migrate_table=plugin_pb2.Write.MessageMigrateTable(table=buf)
+                )
+                record = pa.RecordBatch.from_arrays(
+                    [
+                        pa.array([1, 2, 3]),
+                        pa.array(["a", "b", "c"]),
+                    ],
+                    schema=test_table.to_arrow_schema(),
+                )
+                yield plugin_pb2.Write.Request(
+                    insert=plugin_pb2.Write.MessageInsert(
+                        record=arrow.record_to_bytes(record)
+                    )
+                )
+
+            stub.Write(writer_iterator())
+
+            response = stub.GetTables(plugin_pb2.GetTables.Request(tables=["*"]))
             schemas = arrow.new_schemas_from_bytes(response.tables)
             assert len(schemas) == 1
 
-            response = stub.Sync(plugin_pb2.Sync.Request())
+            response = stub.Sync(plugin_pb2.Sync.Request(tables=["*"]))
             total_records = 0
             for msg in response:
                 if msg.insert is not None: