Working sync proof of concept

Author: hermanschaaf

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

@@ -1,7 +1,8 @@
 import pyarrow as pa
 
 from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc
-from cloudquery.sdk.plugin.plugin import Plugin
+from cloudquery.sdk.message import SyncInsertMessage, SyncMigrateTableMessage
+from cloudquery.sdk.plugin.plugin import Plugin, SyncOptions
 from cloudquery.sdk.schema import tables_to_arrow_schemas
 
 
@@ -33,8 +34,31 @@ def GetTables(self, request: plugin_pb2.GetTables.Request, context):
         return plugin_pb2.GetTables.Response(tables=tablesBytes)
 
     def Sync(self, request, context):
-        plugin_pb2.Sync.Response()
-        return plugin_pb2.Sync.Response()
+        options = SyncOptions(
+            deterministic_cq_id=False,  # TODO
+            skip_dependent_tables=request.skip_dependent_tables,
+            skip_tables=request.skip_tables,
+            tables=request.tables,
+            backend_options=None,
+        )
+
+        for msg in self._plugin.sync(options):
+            if isinstance(msg, SyncInsertMessage):
+                sink = pa.BufferOutputStream()
+                writer = pa.ipc.new_stream(sink, msg.record.schema)
+                writer.write_batch(msg.record)
+                writer.close()
+                buf = sink.getvalue().to_pybytes()
+                yield plugin_pb2.Sync.Response(insert=plugin_pb2.Sync.MessageInsert(
+                    record=buf
+                ))
+            elif isinstance(msg, SyncMigrateTableMessage):
+                yield plugin_pb2.Sync.Response(migrate_table=plugin_pb2.Sync.MessageMigrateTable(
+                    table=msg.table.to_arrow_schema().serialize().to_pybytes()
+                ))
+            else:
+                # unknown sync message type
+                raise NotImplementedError()
 
     def Read(self, request, context):
         raise NotImplementedError()

cloudquery/sdk/message/__init__.py

@@ -1 +1 @@
-from sync import SyncMessage, SyncInsertMessage
+from .sync import SyncMessage, SyncInsertMessage, SyncMigrateTableMessage

cloudquery/sdk/message/sync.py

@@ -1,13 +1,15 @@
-
 import pyarrow as pa
 
+
 class SyncMessage:
     pass
 
+
 class SyncInsertMessage:
     def __init__(self, record: pa.RecordBatch):
-        self._record = record
+        self.record = record
 
-class SyncMigrateMessage:
-    def __init__(self, record: pa.RecordBatch):
-      self._record = record
+
+class SyncMigrateTableMessage:
+    def __init__(self, table: pa.Schema):
+        self.table = table

cloudquery/sdk/message/write.py

@@ -1,13 +1,15 @@
-
 import pyarrow as pa
 
+
 class WriteMessage:
     pass
 
+
 class InsertMessage(WriteMessage):
     def __init__(self, record: pa.RecordBatch):
-        self._record = record
+        self.record = record
+
 
 class MigrateMessage(WriteMessage):
-    def __init__(self, record: pa.RecordBatch):
-      self._record = record
+    def __init__(self, table: pa.Schema):
+        self.table = table

cloudquery/sdk/plugin/plugin.py

@@ -1,7 +1,33 @@
+import queue
+from dataclasses import dataclass
 from typing import List
 
 from cloudquery.sdk.schema import Table
 
+MIGRATE_MODE_STRINGS = ["safe", "force"]
+
+
+@dataclass
+class TableOptions:
+    tables: List[str] = None
+    skip_tables: List[str] = None
+    skip_dependent_tables: bool = False
+
+
+@dataclass
+class BackendOptions:
+    connection: str = None
+    table_name: str = None
+
+
+@dataclass
+class SyncOptions:
+    deterministic_cq_id: bool = False
+    skip_dependent_tables: bool = False
+    skip_tables: List[str] = None
+    tables: List[str] = None
+    backend_options: BackendOptions = None
+
 
 class Plugin:
     def __init__(self, name: str, version: str) -> None:
@@ -17,5 +43,11 @@ def name(self) -> str:
     def version(self) -> str:
         return self._version
 
-    def get_tables(self, tables: List[str], skip_tables: List[str]) -> List[Table]:
+    def get_tables(self, options: TableOptions) -> List[Table]:
+        raise NotImplementedError()
+
+    def sync(self, options: SyncOptions, results: queue.Queue) -> None:
         raise NotImplementedError()
+
+    def close(self) -> None:
+        pass

cloudquery/sdk/schema/__init__.py

@@ -1,3 +1,2 @@
-
 from .column import Column
 from .table import Table, tables_to_arrow_schemas

cloudquery/sdk/schema/table.py

@@ -29,7 +29,7 @@ def primary_keys(self):
     def incremental_keys(self):
         return [column.name for column in self.columns if column.incremental_key]
 
-    def to_arrow_schemas(self):
+    def to_arrow_schema(self):
         fields = []
         md = {
             arrow.METADATA_TABLE_NAME: self.name,
@@ -48,5 +48,5 @@ def __lt__(self, other):
 def tables_to_arrow_schemas(tables: List[Table]):
     schemas = []
     for table in tables:
-        schemas.append(table.to_arrow_schemas())
+        schemas.append(table.to_arrow_schema())
     return schemas

tests/schema/test_table.py

@@ -1,7 +1,8 @@
 import pyarrow as pa
+
 from cloudquery.sdk.schema import Table, Column
 
 
 def test_table():
     table = Table("test_table", [Column("test_column", pa.int32())])
-    table.to_arrow_schemas()
+    table.to_arrow_schema()