-
Notifications
You must be signed in to change notification settings - Fork 2
feat(sync): Send migrate messages #79
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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()), | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| responseObserver); | ||
| } catch (Exception e) { | ||
| responseObserver.onError(e); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,11 @@ | ||
| package io.cloudquery.plugin; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.Getter; | ||
|
|
||
| @AllArgsConstructor | ||
| @Getter | ||
| public class BackendOptions { | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Wrapper class so consumers don't need to mess with the gRPC wrapper |
||
| private final String tableName; | ||
| private final String connection; | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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(); | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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()) { | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Moved from |
||
| 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()); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -88,6 +88,7 @@ private LoggerContext initLogger() { | |
| context.start(configuration); | ||
|
|
||
| logger = context.getLogger(ServeCommand.class.getName()); | ||
| this.plugin.setLogger(logger); | ||
|
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Setting the logger on the plugin was missing |
||
| return context; | ||
| } | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Moved the the
Tableclass under theencodemethod