Skip to content

Commit 085c51f

Browse files
erezrokaherezrokah
authored
feat: Implement getTables (#71)
~~Still doesn't work and I need to handle the options like include, skip, etc.~~ Got the gRPC call to work. Had to: 1. Add `org.apache.arrow:arrow-memory-core` as implementation dependency 2. Add `org.apache.arrow:arrow-memory-nett` as runtime only dependency (see https://mvnrepository.com/artifact/org.apache.arrow/arrow-memory-netty/12.0.1) 3. Add another gradle task so the runtime class is loaded properly. We'll need see how this works with a published version of the SDK.
1 parent a210865 commit 085c51f

File tree

9 files changed

+248
-40
lines changed

9 files changed

+248
-40
lines changed

Diff for: ‎lib/build.gradle

+15-1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@ plugins {
44
id "maven-publish"
55
}
66

7+
ext {
8+
javaMainClass = "io.cloudquery.MainClass"
9+
}
10+
711
group 'io.cloudquery'
812
// x-release-please-start-version
913
version = '0.0.1'
@@ -34,6 +38,7 @@ dependencies {
3438
implementation "io.grpc:grpc-services:1.57.1"
3539
implementation "io.grpc:grpc-testing:1.57.1"
3640
implementation "io.cloudquery:plugin-pb-java:0.0.5"
41+
implementation "org.apache.arrow:arrow-memory-core:12.0.1"
3742
implementation "org.apache.arrow:arrow-vector:12.0.1"
3843

3944
implementation 'org.apache.logging.log4j:log4j-api:2.20.0'
@@ -44,11 +49,12 @@ dependencies {
4449
testImplementation('org.junit.jupiter:junit-jupiter-api:5.10.0')
4550
testImplementation('org.mockito:mockito-core:5.4.0')
4651
testImplementation('org.mockito:mockito-junit-jupiter:5.4.0')
47-
testImplementation('org.apache.arrow:arrow-memory-netty:12.0.1')
4852
testImplementation('nl.jqno.equalsverifier:equalsverifier:3.15')
4953
testRuntimeOnly('org.junit.jupiter:junit-jupiter-engine:5.10.0')
5054

5155
testImplementation 'org.assertj:assertj-core:3.24.2'
56+
57+
runtimeOnly "org.apache.arrow:arrow-memory-netty:12.0.1"
5258
}
5359

5460
test {
@@ -83,3 +89,11 @@ publishing {
8389
}
8490
}
8591
}
92+
93+
task runMemDBServe(type: JavaExec) {
94+
group = "Execution"
95+
description = "Start the MemDB plugin server"
96+
classpath = sourceSets.main.runtimeClasspath
97+
main = javaMainClass
98+
args = ["serve"]
99+
}

Diff for: ‎lib/src/main/java/io/cloudquery/Main.java renamed to ‎lib/src/main/java/io/cloudquery/MainClass.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import io.cloudquery.memdb.MemDB;
44
import io.cloudquery.server.PluginServe;
55

6-
public class Main {
6+
public class MainClass {
77
public static void main(String[] args) {
88
PluginServe serve = PluginServe.builder().plugin(new MemDB()).args(args).build();
99
int exitCode = serve.Serve();

Diff for: ‎lib/src/main/java/io/cloudquery/internal/servers/plugin/v3/PluginServer.java

+112
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,22 @@
11
package io.cloudquery.internal.servers.plugin.v3;
22

33
import io.cloudquery.plugin.v3.PluginGrpc.PluginImplBase;
4+
import io.cloudquery.schema.Table;
5+
import io.cloudquery.plugin.v3.Write;
6+
import io.grpc.stub.StreamObserver;
7+
8+
import java.io.ByteArrayOutputStream;
9+
import java.nio.channels.Channels;
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
13+
import org.apache.arrow.memory.BufferAllocator;
14+
import org.apache.arrow.memory.RootAllocator;
15+
import org.apache.arrow.vector.VectorSchemaRoot;
16+
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
17+
import org.apache.arrow.vector.types.pojo.Schema;
18+
import com.google.protobuf.ByteString;
19+
420
import io.cloudquery.plugin.Plugin;
521

622
public class PluginServer extends PluginImplBase {
@@ -9,4 +25,100 @@ public class PluginServer extends PluginImplBase {
925
public PluginServer(Plugin plugin) {
1026
this.plugin = plugin;
1127
}
28+
29+
@Override
30+
public void getName(io.cloudquery.plugin.v3.GetName.Request request,
31+
StreamObserver<io.cloudquery.plugin.v3.GetName.Response> responseObserver) {
32+
responseObserver
33+
.onNext(io.cloudquery.plugin.v3.GetName.Response.newBuilder().setName(plugin.getName()).build());
34+
responseObserver.onCompleted();
35+
}
36+
37+
@Override
38+
public void getVersion(io.cloudquery.plugin.v3.GetVersion.Request request,
39+
StreamObserver<io.cloudquery.plugin.v3.GetVersion.Response> responseObserver) {
40+
responseObserver.onNext(
41+
io.cloudquery.plugin.v3.GetVersion.Response.newBuilder().setVersion(plugin.getVersion()).build());
42+
responseObserver.onCompleted();
43+
}
44+
45+
@Override
46+
public void init(io.cloudquery.plugin.v3.Init.Request request,
47+
StreamObserver<io.cloudquery.plugin.v3.Init.Response> responseObserver) {
48+
plugin.init();
49+
responseObserver.onNext(io.cloudquery.plugin.v3.Init.Response.newBuilder().build());
50+
responseObserver.onCompleted();
51+
}
52+
53+
@Override
54+
public void getTables(io.cloudquery.plugin.v3.GetTables.Request request,
55+
StreamObserver<io.cloudquery.plugin.v3.GetTables.Response> responseObserver) {
56+
try {
57+
List<Table> tables = plugin.tables();
58+
List<ByteString> byteStrings = new ArrayList<>();
59+
for (Table table : tables) {
60+
try (BufferAllocator bufferAllocator = new RootAllocator()) {
61+
Schema schema = table.toArrowSchema();
62+
VectorSchemaRoot schemaRoot = VectorSchemaRoot.create(schema, bufferAllocator);
63+
try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
64+
try (ArrowStreamWriter writer = new ArrowStreamWriter(schemaRoot, null,
65+
Channels.newChannel(out))) {
66+
writer.start();
67+
writer.end();
68+
byteStrings.add(ByteString.copyFrom(out.toByteArray()));
69+
}
70+
}
71+
}
72+
}
73+
responseObserver
74+
.onNext(io.cloudquery.plugin.v3.GetTables.Response.newBuilder().addAllTables(byteStrings).build());
75+
responseObserver.onCompleted();
76+
} catch (Exception e) {
77+
responseObserver.onError(e);
78+
}
79+
}
80+
81+
@Override
82+
public void sync(io.cloudquery.plugin.v3.Sync.Request request,
83+
StreamObserver<io.cloudquery.plugin.v3.Sync.Response> responseObserver) {
84+
plugin.sync();
85+
responseObserver.onNext(io.cloudquery.plugin.v3.Sync.Response.newBuilder().build());
86+
responseObserver.onCompleted();
87+
}
88+
89+
@Override
90+
public void read(io.cloudquery.plugin.v3.Read.Request request,
91+
StreamObserver<io.cloudquery.plugin.v3.Read.Response> responseObserver) {
92+
plugin.read();
93+
responseObserver.onNext(io.cloudquery.plugin.v3.Read.Response.newBuilder().build());
94+
responseObserver.onCompleted();
95+
}
96+
97+
@Override
98+
public StreamObserver<Write.Request> write(StreamObserver<Write.Response> responseObserver) {
99+
plugin.write();
100+
return new StreamObserver<>() {
101+
@Override
102+
public void onNext(Write.Request request) {
103+
}
104+
105+
@Override
106+
public void onError(Throwable t) {
107+
}
108+
109+
@Override
110+
public void onCompleted() {
111+
responseObserver.onNext(Write.Response.newBuilder().build());
112+
responseObserver.onCompleted();
113+
}
114+
};
115+
}
116+
117+
@Override
118+
public void close(io.cloudquery.plugin.v3.Close.Request request,
119+
StreamObserver<io.cloudquery.plugin.v3.Close.Response> responseObserver) {
120+
plugin.close();
121+
responseObserver.onNext(io.cloudquery.plugin.v3.Close.Response.newBuilder().build());
122+
responseObserver.onCompleted();
123+
}
12124
}

Diff for: ‎lib/src/main/java/io/cloudquery/memdb/MemDB.java

+44
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,53 @@
11
package io.cloudquery.memdb;
22

3+
import java.util.List;
4+
35
import io.cloudquery.plugin.Plugin;
6+
import io.cloudquery.schema.Table;
7+
import io.cloudquery.schema.Column;
8+
import io.cloudquery.schema.SchemaException;
9+
10+
import org.apache.arrow.vector.types.pojo.ArrowType.Utf8;
411

512
public class MemDB extends Plugin {
13+
private List<Table> allTables = List.of(
14+
Table.builder().name("table1").columns(List.of(Column.builder().name("name1").type(new Utf8()).build()))
15+
.build(),
16+
Table.builder().name("table2").columns(List.of(Column.builder().name("name1").type(new Utf8()).build()))
17+
.build());
18+
619
public MemDB() {
720
super("memdb", "0.0.1");
821
}
22+
23+
@Override
24+
public void init() {
25+
// do nothing
26+
}
27+
28+
@Override
29+
public List<Table> tables() throws SchemaException {
30+
return Table.filterDFS(allTables, List.of("*"), List.of(), false);
31+
}
32+
33+
@Override
34+
public void sync() {
35+
// TODO Auto-generated method stub
36+
throw new UnsupportedOperationException("Unimplemented method 'Sync'");
37+
}
38+
39+
@Override
40+
public void read() {
41+
throw new UnsupportedOperationException("Unimplemented method 'Read'");
42+
}
43+
44+
@Override
45+
public void write() {
46+
throw new UnsupportedOperationException("Unimplemented method 'Write'");
47+
}
48+
49+
@Override
50+
public void close() {
51+
// do nothing
52+
}
953
}

Diff for: ‎lib/src/main/java/io/cloudquery/plugin/Plugin.java

+24-2
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,36 @@
11
package io.cloudquery.plugin;
22

3-
import lombok.AllArgsConstructor;
3+
import java.util.List;
4+
5+
import org.apache.logging.log4j.Logger;
6+
7+
import io.cloudquery.schema.SchemaException;
8+
import io.cloudquery.schema.Table;
49
import lombok.Getter;
510
import lombok.NonNull;
11+
import lombok.RequiredArgsConstructor;
12+
import lombok.Setter;
613

14+
@RequiredArgsConstructor
715
@Getter
8-
@AllArgsConstructor
916
public abstract class Plugin {
1017
@NonNull
1118
protected final String name;
1219
@NonNull
1320
protected final String version;
21+
@Setter
22+
protected Logger logger;
23+
24+
public abstract void init();
25+
26+
public abstract List<Table> tables() throws SchemaException;
27+
28+
public abstract void sync();
29+
30+
public abstract void read();
31+
32+
public abstract void write();
33+
34+
public abstract void close();
35+
1436
}

Diff for: ‎lib/src/main/java/io/cloudquery/schema/Resource.java

+3-2
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import io.cloudquery.scalar.ValidationException;
55
import lombok.Builder;
66
import lombok.Getter;
7+
import lombok.NonNull;
78

89
import java.util.ArrayList;
910
import java.util.List;
@@ -17,10 +18,10 @@ public class Resource {
1718
private final List<Scalar<?>> data;
1819

1920
@Builder(toBuilder = true)
20-
public Resource(Table table, Resource parent, Object item) {
21+
public Resource(@NonNull Table table, Resource parent, Object item) {
2122
this.item = item;
2223
this.parent = parent;
23-
this.table = table != null ? table : Table.builder().build();
24+
this.table = table;
2425
this.data = new ArrayList<>();
2526

2627
for (Column column : this.table.getColumns()) {

0 commit comments

Comments
 (0)