Merge pull request #231 from weaviate/grpc-batch-api

Add support for gRPC Batch API
weaviate · Oct 25, 2023 · d202a4a · d202a4a
2 parents f2d09e8 + 479629a
commit d202a4a
Show file tree

Hide file tree

Showing 34 changed files with 56,706 additions and 44 deletions.
diff --git a/pom.xml b/pom.xml
@@ -61,20 +61,52 @@
     <httpclient.version>4.5.14</httpclient.version>
     <lang3.version>3.13.0</lang3.version>
     <junit.version>5.10.0</junit.version>
-    <testcontainers.version>1.19.0</testcontainers.version>
+    <testcontainers.version>1.19.1</testcontainers.version>
     <assertj-core.version>3.24.2</assertj-core.version>
     <jparams.version>1.0.4</jparams.version>
-    <mockito.version>5.5.0</mockito.version>
+    <mockito.version>5.6.0</mockito.version>
     <slf4j.version>2.0.9</slf4j.version>
     <logback.version>1.4.11</logback.version>
     <mock-server.version>5.14.0</mock-server.version>
-    <jackson.version>2.15.2</jackson.version>
-    <oauth2-oidc-sdk.version>10.15</oauth2-oidc-sdk.version>
+    <jackson.version>2.15.3</jackson.version>
+    <oauth2-oidc-sdk.version>11.4</oauth2-oidc-sdk.version>
     <mock-server.version>5.15.0</mock-server.version>
     <jackson.version>2.15.2</jackson.version>
+    <protobuf.java.version>3.24.4</protobuf.java.version>
+    <grpc-netty-shaded.version>1.59.0</grpc-netty-shaded.version>
+    <grpc-protobuf.version>1.59.0</grpc-protobuf.version>
+    <grpc-stub.version>1.59.0</grpc-stub.version>
+    <annotations-api.version>6.0.53</annotations-api.version>
   </properties>
 
   <dependencies>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java</artifactId>
+      <version>${protobuf.java.version}</version>
+    </dependency>
+    <dependency>
+        <groupId>io.grpc</groupId>
+        <artifactId>grpc-netty-shaded</artifactId>
+        <version>${grpc-netty-shaded.version}</version>
+        <scope>runtime</scope>
+      </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-protobuf</artifactId>
+      <version>${grpc-protobuf.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-stub</artifactId>
+      <version>${grpc-stub.version}</version>
+    </dependency>
+    <dependency> <!-- necessary for Java 9+ -->
+      <groupId>org.apache.tomcat</groupId>
+      <artifactId>annotations-api</artifactId>
+      <version>${annotations-api.version}</version>
+      <scope>provided</scope>
+    </dependency>
     <dependency>
       <groupId>org.apache.httpcomponents</groupId>
       <artifactId>httpclient</artifactId>
@@ -157,8 +189,39 @@
   </dependencies>
 
   <build>
+    <extensions>
+      <extension>
+        <groupId>kr.motd.maven</groupId>
+        <artifactId>os-maven-plugin</artifactId>
+        <version>1.7.1</version>
+      </extension>
+    </extensions>
     <pluginManagement>
       <plugins>
+        <plugin>
+          <groupId>org.xolstice.maven.plugins</groupId>
+          <artifactId>protobuf-maven-plugin</artifactId>
+          <version>0.6.1</version>
+          <configuration>
+            <protocArtifact>com.google.protobuf:protoc:3.24.0:exe:${os.detected.classifier}</protocArtifact>
+            <outputDirectory>src/main/java</outputDirectory>
+            <pluginId>grpc-java</pluginId>
+            <pluginArtifact>io.grpc:protoc-gen-grpc-java:1.58.0:exe:${os.detected.classifier}</pluginArtifact>
+            <outputDirectory>src/main/java</outputDirectory>
+            <clearOutputDirectory>false</clearOutputDirectory>
+
+          </configuration>
+          <executions>
+            <execution>
+              <goals>
+                <goal>compile</goal>
+                <goal>compile-custom</goal>
+                <goal>test-compile</goal>
+                <goal>test-compile-custom</goal>
+              </goals>
+            </execution>
+          </executions>
+        </plugin>
         <plugin>
           <artifactId>maven-clean-plugin</artifactId>
           <version>3.1.0</version>
@@ -344,7 +407,7 @@
         </plugin>
         <plugin>
           <artifactId>maven-project-info-reports-plugin</artifactId>
-          <version>3.1.1</version>
+          <version>3.1.1my</version>
         </plugin>
       </plugins>
     </pluginManagement>
@@ -353,6 +416,10 @@
         <groupId>org.projectlombok</groupId>
         <artifactId>lombok-maven-plugin</artifactId>
       </plugin>
+      <plugin>
+        <groupId>org.xolstice.maven.plugins</groupId>
+        <artifactId>protobuf-maven-plugin</artifactId>
+      </plugin>
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-antrun-plugin</artifactId>

diff --git a/src/main/java/io/weaviate/client/Config.java b/src/main/java/io/weaviate/client/Config.java
@@ -1,21 +1,35 @@
 package io.weaviate.client;
 
 import java.util.Map;
+import lombok.Getter;
+import lombok.Setter;
 
 public class Config {
 
   private static final int DEFAULT_TIMEOUT_SECONDS = 60;
+  @Getter
   private final String scheme;
+  @Getter
   private final String host;
   private final String version;
+  @Getter
   private final Map<String, String> headers;
+  @Getter
   private final int connectionTimeout;
+  @Getter
   private final int connectionRequestTimeout;
+  @Getter
   private final int socketTimeout;
+  @Getter
   private String proxyHost;
+  @Getter
   private int proxyPort;
+  @Getter
   private String proxyScheme;
-
+  @Setter
+  private boolean useGRPC;
+  @Getter @Setter
+  private String grpcAddress;
 
   public Config(String scheme, String host) {
     this(scheme, host, null, DEFAULT_TIMEOUT_SECONDS, DEFAULT_TIMEOUT_SECONDS, DEFAULT_TIMEOUT_SECONDS);
@@ -35,24 +49,19 @@ public Config(String scheme, String host, Map<String, String> headers, int conne
     this.socketTimeout = socketTimeout;
   }
 
-  public String getBaseURL() {
-    return scheme + "://" + host + "/" + version;
-  }
-
-  public Map<String, String> getHeaders() {
-    return headers;
-  }
-
-  public int getConnectionTimeout() {
-    return connectionTimeout;
-  }
-
-  public int getConnectionRequestTimeout() {
-    return connectionRequestTimeout;
+  public Config(String scheme, String host, Map<String, String> headers, int timeout, boolean useGRPC) {
+    this.scheme = scheme;
+    this.host = host;
+    this.version = "v1";
+    this.headers = headers;
+    this.connectionTimeout = timeout;
+    this.connectionRequestTimeout = timeout;
+    this.socketTimeout = timeout;
+    this.useGRPC = useGRPC;
   }
 
-  public int getSocketTimeout() {
-    return socketTimeout;
+  public String getBaseURL() {
+    return scheme + "://" + host + "/" + version;
   }
 
   public void setProxy(String proxyHost, int proxyPort, String proxyScheme) {
@@ -61,16 +70,7 @@ public void setProxy(String proxyHost, int proxyPort, String proxyScheme) {
     this.proxyScheme = proxyScheme;
   }
 
-  public String getProxyHost() {
-    return proxyHost;
-  }
-
-  public int getProxyPort() {
-    return proxyPort;
-  }
-
-  public String getProxyScheme() {
-    return proxyScheme;
+  public boolean useGRPC() {
+    return this.useGRPC;
   }
-
 }
diff --git a/src/main/java/io/weaviate/client/base/grpc/GrpcClient.java b/src/main/java/io/weaviate/client/base/grpc/GrpcClient.java
@@ -0,0 +1,44 @@
+package io.weaviate.client.base.grpc;
+
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import io.grpc.Metadata;
+import io.grpc.stub.MetadataUtils;
+import io.weaviate.client.Config;
+import io.weaviate.client.grpc.protocol.v1.WeaviateGrpc;
+import java.util.Map;
+
+public class GrpcClient {
+
+  public static WeaviateGrpc.WeaviateBlockingStub create(Config config) {
+    Metadata headers = new Metadata();
+    if (config.getHeaders() != null) {
+      for (Map.Entry<String, String> e : config.getHeaders().entrySet()) {
+        headers.put(Metadata.Key.of(e.getKey(), Metadata.ASCII_STRING_MARSHALLER), e.getValue());
+      }
+    }
+    ManagedChannelBuilder<?> channelBuilder = ManagedChannelBuilder.forTarget(getAddress(config));
+    if (config.getScheme().equals("https")) {
+      channelBuilder = channelBuilder.useTransportSecurity();
+    } else {
+      channelBuilder.usePlaintext();
+    }
+    ManagedChannel channel = channelBuilder.build();
+    WeaviateGrpc.WeaviateBlockingStub blockingStub = WeaviateGrpc.newBlockingStub(channel);
+    return blockingStub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
+  }
+
+  private static String getAddress(Config config) {
+    if (config.getGrpcAddress() != null) {
+      return config.getGrpcAddress();
+    }
+    String host = config.getHost();
+    if (!host.contains(":")) {
+      if (config.getScheme() != null && config.getScheme().equals("https")) {
+        return String.format("%s:443", host);
+      }
+      return String.format("%s:80", host);
+    }
+    return host;
+  }
+}