feat(sync): Send migrate messages (#79)

First part of the sync, sending migrate table messages to the client

Author: erezrokah

lib/src/main/java/io/cloudquery/internal/servers/plugin/v3/PluginServer.java

@@ -1,20 +1,14 @@
 package io.cloudquery.internal.servers.plugin.v3;
 
 import com.google.protobuf.ByteString;
+import io.cloudquery.plugin.BackendOptions;
 import io.cloudquery.plugin.Plugin;
 import io.cloudquery.plugin.v3.PluginGrpc.PluginImplBase;
 import io.cloudquery.plugin.v3.Write;
 import io.cloudquery.schema.Table;
 import io.grpc.stub.StreamObserver;
-import java.io.ByteArrayOutputStream;
-import java.nio.channels.Channels;
 import java.util.ArrayList;
 import java.util.List;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.RootAllocator;
-import org.apache.arrow.vector.VectorSchemaRoot;
-import org.apache.arrow.vector.ipc.ArrowStreamWriter;
-import org.apache.arrow.vector.types.pojo.Schema;
 
 public class PluginServer extends PluginImplBase {
   private final Plugin plugin;
@@ -64,18 +58,7 @@ public void getTables(
               request.getSkipDependentTables());
       List<ByteString> byteStrings = new ArrayList<>();
       for (Table table : tables) {
-        try (BufferAllocator bufferAllocator = new RootAllocator()) {
-          Schema schema = table.toArrowSchema();
-          VectorSchemaRoot schemaRoot = VectorSchemaRoot.create(schema, bufferAllocator);
-          try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
-            try (ArrowStreamWriter writer =
-                new ArrowStreamWriter(schemaRoot, null, Channels.newChannel(out))) {
-              writer.start();
-              writer.end();
-              byteStrings.add(ByteString.copyFrom(out.toByteArray()));
-            }
-          }
-        }
+        byteStrings.add(table.encode());
       }
       responseObserver.onNext(
           io.cloudquery.plugin.v3.GetTables.Response.newBuilder()
@@ -91,9 +74,18 @@ public void getTables(
   public void sync(
       io.cloudquery.plugin.v3.Sync.Request request,
       StreamObserver<io.cloudquery.plugin.v3.Sync.Response> responseObserver) {
-    plugin.sync();
-    responseObserver.onNext(io.cloudquery.plugin.v3.Sync.Response.newBuilder().build());
-    responseObserver.onCompleted();
+    try {
+      plugin.sync(
+          request.getTablesList(),
+          request.getSkipTablesList(),
+          request.getSkipDependentTables(),
+          request.getDeterministicCqId(),
+          new BackendOptions(
+              request.getBackend().getTableName(), request.getBackend().getConnection()),
+          responseObserver);
+    } catch (Exception e) {
+      responseObserver.onError(e);
+    }
   }
 
   @Override

lib/src/main/java/io/cloudquery/memdb/MemDB.java

@@ -1,9 +1,12 @@
 package io.cloudquery.memdb;
 
+import io.cloudquery.plugin.BackendOptions;
 import io.cloudquery.plugin.Plugin;
+import io.cloudquery.scheduler.Scheduler;
 import io.cloudquery.schema.Column;
 import io.cloudquery.schema.SchemaException;
 import io.cloudquery.schema.Table;
+import io.grpc.stub.StreamObserver;
 import java.util.List;
 import org.apache.arrow.vector.types.pojo.ArrowType.Utf8;
 
@@ -36,9 +39,22 @@ public List<Table> tables(
   }
 
   @Override
-  public void sync() {
-    // TODO Auto-generated method stub
-    throw new UnsupportedOperationException("Unimplemented method 'Sync'");
+  public void sync(
+      List<String> includeList,
+      List<String> skipList,
+      boolean skipDependentTables,
+      boolean deterministicCqId,
+      BackendOptions backendOptions,
+      StreamObserver<io.cloudquery.plugin.v3.Sync.Response> syncStream)
+      throws SchemaException {
+    List<Table> filtered = Table.filterDFS(allTables, includeList, skipList, skipDependentTables);
+    Scheduler.builder()
+        .tables(filtered)
+        .syncStream(syncStream)
+        .deterministicCqId(deterministicCqId)
+        .logger(getLogger())
+        .build()
+        .sync();
   }
 
   @Override

lib/src/main/java/io/cloudquery/plugin/BackendOptions.java

@@ -0,0 +1,11 @@
+package io.cloudquery.plugin;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+
+@AllArgsConstructor
+@Getter
+public class BackendOptions {
+  private final String tableName;
+  private final String connection;
+}

lib/src/main/java/io/cloudquery/plugin/Plugin.java

@@ -2,6 +2,7 @@
 
 import io.cloudquery.schema.SchemaException;
 import io.cloudquery.schema.Table;
+import io.grpc.stub.StreamObserver;
 import java.util.List;
 import lombok.Getter;
 import lombok.NonNull;
@@ -22,7 +23,14 @@ public abstract List<Table> tables(
       List<String> includeList, List<String> skipList, boolean skipDependentTables)
       throws SchemaException;
 
-  public abstract void sync();
+  public abstract void sync(
+      List<String> includeList,
+      List<String> skipList,
+      boolean skipDependentTables,
+      boolean deterministicCqId,
+      BackendOptions backendOptions,
+      StreamObserver<io.cloudquery.plugin.v3.Sync.Response> syncStream)
+      throws SchemaException;
 
   public abstract void read();
 

lib/src/main/java/io/cloudquery/scheduler/Scheduler.java

@@ -1,5 +1,35 @@
 package io.cloudquery.scheduler;
 
+import io.cloudquery.plugin.v3.Sync;
+import io.cloudquery.schema.Table;
+import io.grpc.stub.StreamObserver;
+import java.io.IOException;
+import java.util.List;
+import lombok.Builder;
+import lombok.NonNull;
+import org.apache.logging.log4j.Logger;
+
+@Builder
 public class Scheduler {
-  public Scheduler() {}
+  @NonNull private final List<Table> tables;
+  @NonNull private final StreamObserver<io.cloudquery.plugin.v3.Sync.Response> syncStream;
+  @NonNull private final Logger logger;
+
+  private boolean deterministicCqId;
+
+  public void sync() {
+    for (Table table : tables) {
+      try {
+        logger.info("sending migrate message for table: {}", table.getName());
+        Sync.MessageMigrateTable migrateTable =
+            Sync.MessageMigrateTable.newBuilder().setTable(table.encode()).build();
+        Sync.Response response = Sync.Response.newBuilder().setMigrateTable(migrateTable).build();
+        syncStream.onNext(response);
+      } catch (IOException e) {
+        syncStream.onError(e);
+        return;
+      }
+    }
+    syncStream.onCompleted();
+  }
 }

lib/src/main/java/io/cloudquery/schema/Table.java

@@ -2,9 +2,13 @@
 
 import static java.util.Arrays.asList;
 
+import com.google.protobuf.ByteString;
 import io.cloudquery.glob.Glob;
 import io.cloudquery.schema.Column.ColumnBuilder;
 import io.cloudquery.transformers.TransformerException;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.channels.Channels;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -16,6 +20,10 @@
 import lombok.Getter;
 import lombok.NonNull;
 import lombok.Setter;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ipc.ArrowStreamWriter;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
 
@@ -228,4 +236,19 @@ public Schema toArrowSchema() {
     Schema schema = new Schema(asList(fields), metadata);
     return schema;
   }
+
+  public ByteString encode() throws IOException {
+    try (BufferAllocator bufferAllocator = new RootAllocator()) {
+      Schema schema = toArrowSchema();
+      VectorSchemaRoot schemaRoot = VectorSchemaRoot.create(schema, bufferAllocator);
+      try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+        try (ArrowStreamWriter writer =
+            new ArrowStreamWriter(schemaRoot, null, Channels.newChannel(out))) {
+          writer.start();
+          writer.end();
+          return ByteString.copyFrom(out.toByteArray());
+        }
+      }
+    }
+  }
 }

lib/src/main/java/io/cloudquery/server/ServeCommand.java

@@ -88,6 +88,7 @@ private LoggerContext initLogger() {
     context.start(configuration);
 
     logger = context.getLogger(ServeCommand.class.getName());
+    this.plugin.setLogger(logger);
     return context;
   }