diff --git a/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoAdminClient.java b/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoAdminClient.java index 3429414c521..59393e8f58c 100644 --- a/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoAdminClient.java +++ b/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoAdminClient.java @@ -40,9 +40,11 @@ public class GravitinoAdminClient extends GravitinoClientBase implements Support * * @param uri The base URI for the Gravitino API. * @param authDataProvider The provider of the data which is used for authentication. + * @param headers The base header for Gravitino API. */ - private GravitinoAdminClient(String uri, AuthDataProvider authDataProvider) { - super(uri, authDataProvider); + private GravitinoAdminClient( + String uri, AuthDataProvider authDataProvider, Map headers) { + super(uri, authDataProvider, headers); } /** @@ -188,7 +190,7 @@ public GravitinoAdminClient build() { Preconditions.checkArgument( uri != null && !uri.isEmpty(), "The argument 'uri' must be a valid URI"); - return new GravitinoAdminClient(uri, authDataProvider); + return new GravitinoAdminClient(uri, authDataProvider, headers); } } } diff --git a/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClient.java b/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClient.java index dd391266b6d..f5584adb625 100644 --- a/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClient.java +++ b/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClient.java @@ -33,10 +33,15 @@ public class GravitinoClient extends GravitinoClientBase implements SupportsCata * @param uri The base URI for the Gravitino API. * @param metalakeName The specified metalake name. * @param authDataProvider The provider of the data which is used for authentication. + * @param headers The base header for Gravitino API. * @throws NoSuchMetalakeException if the metalake with specified name does not exist. */ - private GravitinoClient(String uri, String metalakeName, AuthDataProvider authDataProvider) { - super(uri, authDataProvider); + private GravitinoClient( + String uri, + String metalakeName, + AuthDataProvider authDataProvider, + Map headers) { + super(uri, authDataProvider, headers); this.metalake = loadMetalake(NameIdentifier.of(metalakeName)); } @@ -138,7 +143,7 @@ public GravitinoClient build() { metalakeName != null && !metalakeName.isEmpty(), "The argument 'metalakeName' must be a valid name"); - return new GravitinoClient(uri, metalakeName, authDataProvider); + return new GravitinoClient(uri, metalakeName, authDataProvider, headers); } } } diff --git a/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClientBase.java b/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClientBase.java index 75350a7b954..87a4410a4a3 100644 --- a/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClientBase.java +++ b/clients/client-java/src/main/java/com/datastrato/gravitino/client/GravitinoClientBase.java @@ -9,10 +9,12 @@ import com.datastrato.gravitino.dto.responses.MetalakeResponse; import com.datastrato.gravitino.dto.responses.VersionResponse; import com.datastrato.gravitino.exceptions.NoSuchMetalakeException; +import com.google.common.collect.ImmutableMap; import java.io.Closeable; import java.net.URI; import java.net.URISyntaxException; import java.util.Collections; +import java.util.Map; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -37,12 +39,15 @@ public abstract class GravitinoClientBase implements Closeable { * * @param uri The base URI for the Gravitino API. * @param authDataProvider The provider of the data which is used for authentication. + * @param headers The base header of the Gravitino API. */ - protected GravitinoClientBase(String uri, AuthDataProvider authDataProvider) { + protected GravitinoClientBase( + String uri, AuthDataProvider authDataProvider, Map headers) { this.restClient = HTTPClient.builder(Collections.emptyMap()) .uri(uri) .withAuthDataProvider(authDataProvider) + .withHeaders(headers) .build(); } @@ -103,6 +108,8 @@ public abstract static class Builder { protected String uri; /** The authentication provider. */ protected AuthDataProvider authDataProvider; + /** The request base header for the Gravitino API. */ + protected Map headers = ImmutableMap.of(); /** * The constructor for the Builder class. @@ -154,6 +161,19 @@ public Builder withKerberosAuth(KerberosTokenProvider dataProvider) { return this; } + /** + * Set base header for Gravitino Client. + * + * @param headers the base header. + * @return This Builder instance for method chaining. + */ + public Builder withHeaders(Map headers) { + if (headers != null) { + this.headers = ImmutableMap.copyOf(headers); + } + return this; + } + /** * Builds a new instance. Subclasses should overwrite this method. * diff --git a/clients/client-java/src/test/java/com/datastrato/gravitino/client/TestGravitinoClientBuilder.java b/clients/client-java/src/test/java/com/datastrato/gravitino/client/TestGravitinoClientBuilder.java new file mode 100644 index 00000000000..965632e64fa --- /dev/null +++ b/clients/client-java/src/test/java/com/datastrato/gravitino/client/TestGravitinoClientBuilder.java @@ -0,0 +1,79 @@ +/* + * Copyright 2024 Datastrato Pvt Ltd. + * This software is licensed under the Apache License version 2. + */ +package com.datastrato.gravitino.client; + +import com.google.common.collect.ImmutableMap; +import java.util.Map; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestGravitinoClientBuilder { + @Test + public void testGravitinoClientHeaders() { + Map headers = ImmutableMap.of("k1", "v1"); + try (MockGravitinoClient client = + MockGravitinoClient.builder("http://127.0.0.1").withHeaders(headers).build()) { + Assertions.assertEquals(headers, client.getHeaders()); + } + + try (MockGravitinoClient client1 = MockGravitinoClient.builder("http://127.0.0.1").build()) { + Assertions.assertEquals(ImmutableMap.of(), client1.getHeaders()); + } + + try (MockGravitinoClient client1 = + MockGravitinoClient.builder("http://127.0.0.1").withHeaders(null).build()) { + Assertions.assertEquals(ImmutableMap.of(), client1.getHeaders()); + } + } +} + +class MockGravitinoClient extends GravitinoClientBase { + + private Map headers; + + /** + * Constructs a new GravitinoClient with the given URI, authenticator and AuthDataProvider. + * + * @param uri The base URI for the Gravitino API. + * @param authDataProvider The provider of the data which is used for authentication. + * @param headers The base header of the Gravitino API. + */ + private MockGravitinoClient( + String uri, AuthDataProvider authDataProvider, Map headers) { + super(uri, authDataProvider, headers); + this.headers = headers; + } + + Map getHeaders() { + return headers; + } + + /** + * Creates a new builder for constructing a GravitinoClient. + * + * @param uri The base URI for the Gravitino API. + * @return A new instance of the Builder class for constructing a GravitinoClient. + */ + static MockGravitinoClientBuilder builder(String uri) { + return new MockGravitinoClientBuilder(uri); + } + + static class MockGravitinoClientBuilder extends GravitinoClientBase.Builder { + + /** + * The constructor for the Builder class. + * + * @param uri The base URI for the Gravitino API. + */ + protected MockGravitinoClientBuilder(String uri) { + super(uri); + } + + @Override + public MockGravitinoClient build() { + return new MockGravitinoClient(uri, authDataProvider, headers); + } + } +}