Merge pull request #106 from qdrant/custom-headers

feat: Add support for custom headers

Author: generall

src/main/java/io/qdrant/client/ApiKeyCredentials.java

@@ -0,0 +1,48 @@
+package io.qdrant.client;
+
+import io.grpc.CallCredentials;
+import io.grpc.Metadata;
+import io.grpc.Status;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.Executor;
+import javax.annotation.Nullable;
+
+/** Used internally by the client to send the API key and any headers as gRPC metadata. */
+public class MetadataCredentials extends CallCredentials {
+  @Nullable private final String apiKey;
+  private final Map<String, String> headers;
+
+  /**
+   * Instantiates a new instance of {@link MetadataCredentials}
+   *
+   * @param apiKey The API key to use for authentication.
+   * @param headers Custom headers to send with every request.
+   */
+  public MetadataCredentials(@Nullable String apiKey, Map<String, String> headers) {
+    this.apiKey = apiKey;
+    this.headers = headers != null ? headers : Collections.emptyMap();
+  }
+
+  @Override
+  public void applyRequestMetadata(
+      RequestInfo requestInfo, Executor appExecutor, MetadataApplier applier) {
+    appExecutor.execute(
+        () -> {
+          try {
+            Metadata metadata = new Metadata();
+            if (apiKey != null) {
+              metadata.put(Metadata.Key.of("api-key", Metadata.ASCII_STRING_MARSHALLER), apiKey);
+            }
+            for (Map.Entry<String, String> entry : headers.entrySet()) {
+              metadata.put(
+                  Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER),
+                  entry.getValue());
+            }
+            applier.apply(metadata);
+          } catch (Throwable e) {
+            applier.fail(Status.INTERNAL.withCause(e));
+          }
+        });
+  }
+}

src/main/java/io/qdrant/client/MetadataCredentials.java

@@ -10,6 +10,7 @@
 import io.qdrant.client.grpc.QdrantGrpc.QdrantFutureStub;
 import io.qdrant.client.grpc.SnapshotsGrpc.SnapshotsFutureStub;
 import java.time.Duration;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
@@ -193,8 +194,10 @@ public static class Builder {
     private final ManagedChannel channel;
     private final boolean shutdownChannelOnClose;
     private final boolean checkCompatibility;
+    @Nullable private String apiKey;
     @Nullable private CallCredentials callCredentials;
     @Nullable private Duration timeout;
+    @Nullable private Map<String, String> headers;
 
     Builder(ManagedChannel channel, boolean shutdownChannelOnClose, boolean checkCompatibility) {
       this.channel = channel;
@@ -218,7 +221,7 @@ public static class Builder {
      * @return this
      */
     public Builder withApiKey(String apiKey) {
-      this.callCredentials = new ApiKeyCredentials(apiKey);
+      this.apiKey = apiKey;
       return this;
     }
 
@@ -237,6 +240,9 @@ public Builder withTimeout(@Nullable Duration timeout) {
      * Sets the credential data that will be propagated to the server via request metadata for each
      * RPC.
      *
+     * <p>Note: If both {@link #withApiKey(String)} / {@link #withHeaders(Map)} and this method are
+     * used, this method takes precedence and the API key and headers will be ignored.
+     *
      * @param callCredentials The call credentials to use.
      * @return this
      */
@@ -245,6 +251,17 @@ public Builder withCallCredentials(@Nullable CallCredentials callCredentials) {
       return this;
     }
 
+    /**
+     * Sets custom headers to send with every gRPC request.
+     *
+     * @param headers The headers to send.
+     * @return this
+     */
+    public Builder withHeaders(Map<String, String> headers) {
+      this.headers = headers;
+      return this;
+    }
+
     /**
      * Builds a new instance of {@link QdrantGrpcClient}
      *
@@ -255,7 +272,13 @@ public QdrantGrpcClient build() {
         String clientVersion = Builder.class.getPackage().getImplementationVersion();
         checkVersionsCompatibility(clientVersion);
       }
-      return new QdrantGrpcClient(channel, shutdownChannelOnClose, callCredentials, timeout);
+
+      CallCredentials credentials = this.callCredentials;
+      if (credentials == null && (apiKey != null || headers != null)) {
+        credentials = new MetadataCredentials(apiKey, headers);
+      }
+
+      return new QdrantGrpcClient(channel, shutdownChannelOnClose, credentials, timeout);
     }
 
     private static ManagedChannel createChannel(