feat: add http customizer to client builder (#2822)

Signed-off-by: David Goss <[email protected]>
Co-authored-by: Willy Lulciuc <[email protected]>

Author: davidjgoss

clients/java/README.md

@@ -27,36 +27,32 @@ implementation 'io.github.marquezproject:marquez-java:0.47.0
 // Connect to http://localhost:5000
 MarquezClient client = MarquezClient.builder()
   .baseUrl("http://localhost:5000")
-  .build()
+  .build();
 
 // List namespaces
 List<Namespace> namespaces = client.listNamespaces();
 ```
-### Writing Metadata
-To collect OpenLineage events using Marquez, please use the [openlineage-java](https://search.maven.org/artifact/io.openlineage/openlineage-java) library. OpenLineage is an Open Standard for lineage metadata collection designed to collect metadata for a job in execution.
 
-## HTTPS
+### Customization
 
-```java
-import javax.net.ssl.KeyManager;
-import javax.net.ssl.SSLContext;
-import javax.net.ssl.TrustManager;
-import marquez.client.MarquezClient;
-.
-.
-KeyManager[] keyManager = setUpKeyManagers();
-TrustManager[] trustManager = setUpTrustManagers();
+The client uses Apache's `httpclient` under the hood. The defaults can be a bit limiting, but you can optionally provide a function that accepts the `HttpClientBuilder` so you can customize it before it's finalised:
 
+```java
 SSLContext sslContext = SSLContext.getInstance("TLS");
-sslContext.init(keyManager, trustManager, null);
+HttpClientConnectionManager connectionManager = new PoolingHttpClientConnectionManager(...);
 
-// Connect to https://localhost:5000
 MarquezClient client = MarquezClient.builder()
-  .sslContext(sslContext)
-  .baseUrl("https://localhost:5000")
+  .baseUrl("http://localhost:5000")
+  .httpCustomizer(httpClientBuilder -> {
+    httpClientBuilder.setSSLContext(sslContext);
+    httpClientBuilder.setConnectionManager(connectionManager);
+  })
   .build();
 ```
 
+### Writing Metadata
+To collect OpenLineage events using Marquez, please use the [openlineage-java](https://search.maven.org/artifact/io.openlineage/openlineage-java) library. OpenLineage is an Open Standard for lineage metadata collection designed to collect metadata for a job in execution.
+
 ----
 SPDX-License-Identifier: Apache-2.0
 Copyright 2018-2023 contributors to the Marquez project.

clients/java/src/main/java/marquez/client/MarquezClient.java

@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Properties;
 import java.util.Set;
+import java.util.function.Consumer;
 import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
 import lombok.AllArgsConstructor;
@@ -54,6 +55,7 @@
 import marquez.client.models.Source;
 import marquez.client.models.SourceMeta;
 import marquez.client.models.Tag;
+import org.apache.http.impl.client.HttpClientBuilder;
 
 @Slf4j
 public class MarquezClient {
@@ -506,6 +508,7 @@ public static final class Builder {
     @VisibleForTesting URL baseUrl;
     @VisibleForTesting @Nullable String apiKey;
     @VisibleForTesting @Nullable SSLContext sslContext;
+    @VisibleForTesting @Nullable Consumer<HttpClientBuilder> httpCustomizer;
 
     private Builder() {
       this.baseUrl = DEFAULT_BASE_URL;
@@ -530,10 +533,15 @@ public Builder sslContext(@Nullable SSLContext sslContext) {
       return this;
     }
 
+    public Builder customize(@Nullable Consumer<HttpClientBuilder> httpCustomizer) {
+      this.httpCustomizer = httpCustomizer;
+      return this;
+    }
+
     public MarquezClient build() {
       return new MarquezClient(
           MarquezUrl.create(baseUrl),
-          MarquezHttp.create(sslContext, MarquezClient.Version.get(), apiKey));
+          MarquezHttp.create(sslContext, httpCustomizer, MarquezClient.Version.get(), apiKey));
     }
   }
 

clients/java/src/main/java/marquez/client/MarquezHttp.java

@@ -17,6 +17,7 @@
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.net.URL;
+import java.util.function.Consumer;
 import javax.annotation.Nullable;
 import javax.net.ssl.SSLContext;
 import lombok.Getter;
@@ -59,15 +60,17 @@ static MarquezHttp create(
 
   static MarquezHttp create(
       @Nullable SSLContext sslContext,
+      @Nullable Consumer<HttpClientBuilder> httpCustomizer,
       @NonNull final MarquezClient.Version version,
       @Nullable final String apiKey) {
     final UserAgent userAgent = UserAgent.of(version);
-    final CloseableHttpClient http =
-        HttpClientBuilder.create()
-            .setUserAgent(userAgent.getValue())
-            .setSSLContext(sslContext)
-            .build();
-    return new MarquezHttp(http, apiKey);
+    final HttpClientBuilder httpClientBuilder = HttpClientBuilder.create();
+    httpClientBuilder.setUserAgent(userAgent.getValue());
+    httpClientBuilder.setSSLContext(sslContext);
+    if (httpCustomizer != null) {
+      httpCustomizer.accept(httpClientBuilder);
+    }
+    return new MarquezHttp(httpClientBuilder.build(), apiKey);
   }
 
   String post(URL url) {

clients/java/src/test/java/marquez/client/MarquezClientTest.java

@@ -547,6 +547,17 @@ public void testClientBuilder_sslContext()
     builder.build();
   }
 
+  @Test
+  public void testClientBuilder_httpCustomizer() {
+    MarquezClient.Builder builder = MarquezClient.builder();
+    assertThat(builder.httpCustomizer == null);
+
+    builder.customize(httpClientBuilder -> httpClientBuilder.setMaxConnTotal(30));
+    assertThat(builder.httpCustomizer != null);
+
+    builder.build();
+  }
+
   @Test
   public void testCreateNamespace() throws Exception {
     final URL url = buildUrlFor("/namespaces/%s", NAMESPACE_NAME);

clients/java/src/test/java/marquez/client/MarquezHttpTest.java

@@ -85,9 +85,9 @@ public void testNewMarquezHttp_throwsOnNull() {
     assertThatNullPointerException().isThrownBy(() -> MarquezHttp.create(null));
     assertThatNullPointerException().isThrownBy(() -> MarquezHttp.create(null, API_KEY));
     assertThatNullPointerException().isThrownBy(() -> MarquezHttp.create(null, null));
-    assertThatNullPointerException().isThrownBy(() -> MarquezHttp.create(null, null, null));
+    assertThatNullPointerException().isThrownBy(() -> MarquezHttp.create(null, null, null, null));
     assertThatNullPointerException()
-        .isThrownBy(() -> MarquezHttp.create(SSLContext.getDefault(), null, API_KEY));
+        .isThrownBy(() -> MarquezHttp.create(SSLContext.getDefault(), null, null, API_KEY));
   }
 
   @Test