Merge pull request #301 from databricks/bhu1rdb/sslClient

Implemented makeClosableDisabledSslHttpClient for exporting metrics

Author: bhu1rdb

src/main/java/com/databricks/jdbc/client/http/DatabricksHttpClient.java

@@ -16,10 +16,14 @@
 import com.databricks.sdk.core.utils.ProxyUtils;
 import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
+import java.security.KeyManagementException;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import javax.net.ssl.SSLContext;
 import org.apache.http.HttpException;
 import org.apache.http.HttpHost;
 import org.apache.http.HttpResponse;
@@ -28,11 +32,13 @@
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.conn.UnsupportedSchemeException;
 import org.apache.http.conn.routing.HttpRoute;
+import org.apache.http.conn.ssl.NoopHostnameVerifier;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.conn.DefaultSchemePortResolver;
 import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.apache.http.protocol.HttpContext;
+import org.apache.http.ssl.SSLContextBuilder;
 
 /** Http client implementation to be used for executing http requests. */
 public class DatabricksHttpClient implements IDatabricksHttpClient {
@@ -72,6 +78,7 @@ public class DatabricksHttpClient implements IDatabricksHttpClient {
   private static boolean shouldRetryRateLimitError;
   private static int rateLimitRetryTimeout;
   protected static int idleHttpConnectionExpiry;
+  private CloseableHttpClient httpDisabledSSLClient;
 
   private DatabricksHttpClient(IDatabricksConnectionContext connectionContext) {
     initializeConnectionManager();
@@ -81,6 +88,7 @@ private DatabricksHttpClient(IDatabricksConnectionContext connectionContext) {
     shouldRetryRateLimitError = connectionContext.shouldRetryRateLimitError();
     rateLimitRetryTimeout = connectionContext.getRateLimitRetryTimeout();
     httpClient = makeClosableHttpClient(connectionContext);
+    httpDisabledSSLClient = makeClosableDisabledSslHttpClient();
     idleHttpConnectionExpiry = connectionContext.getIdleHttpConnectionExpiry();
   }
 
@@ -143,6 +151,26 @@ private CloseableHttpClient makeClosableHttpClient(
     return builder.build();
   }
 
+  private CloseableHttpClient makeClosableDisabledSslHttpClient() {
+    try {
+      // Create SSL context that trusts all certificates
+      SSLContext sslContext =
+          new SSLContextBuilder().loadTrustMaterial(null, (chain, authType) -> true).build();
+
+      // Create HttpClient with the SSL context
+      return HttpClientBuilder.create()
+          .setSSLContext(sslContext)
+          .setSSLHostnameVerifier(new NoopHostnameVerifier())
+          .build();
+    } catch (NoSuchAlgorithmException | KeyManagementException | KeyStoreException e) {
+      LoggingUtil.log(
+          LogLevel.DEBUG,
+          String.format(
+              "Error in creating HttpClient with the SSL context [{%s}]", e.getMessage()));
+    }
+    return null;
+  }
+
   private boolean handleRetry(IOException exception, int executionCount, HttpContext context) {
     int errCode = getErrorCode(exception);
     if (!isErrorCodeRetryable(errCode)) {
@@ -322,20 +350,22 @@ public CloseableHttpResponse execute(HttpUriRequest request) throws DatabricksHt
     try {
       return httpClient.execute(request);
     } catch (IOException e) {
-      Throwable cause = e;
-      while (cause != null) {
-        if (cause instanceof DatabricksRetryHandlerException) {
-          throw new DatabricksHttpException(cause.getMessage(), cause);
-        }
-        cause = cause.getCause();
-      }
-      String errorMsg =
-          String.format(
-              "Caught error while executing http request: [%s]. Error Message: [%s]",
-              RequestSanitizer.sanitizeRequest(request), e);
-      LoggingUtil.log(LogLevel.ERROR, errorMsg);
-      throw new DatabricksHttpException(errorMsg, e);
+      throwHttpException(e, request, LogLevel.ERROR);
+    }
+    return null;
+  }
+
+  public CloseableHttpResponse executeWithoutSSL(HttpUriRequest request)
+      throws DatabricksHttpException {
+    LoggingUtil.log(
+        LogLevel.DEBUG,
+        String.format("Executing HTTP request [{%s}]", RequestSanitizer.sanitizeRequest(request)));
+    try {
+      return httpDisabledSSLClient.execute(request);
+    } catch (Exception e) {
+      throwHttpException(e, request, LogLevel.DEBUG);
     }
+    return null;
   }
 
   @Override
@@ -385,4 +415,21 @@ public static synchronized void removeInstance(IDatabricksConnectionContext cont
       }
     }
   }
+
+  private static void throwHttpException(Exception e, HttpUriRequest request, LogLevel logLevel)
+      throws DatabricksHttpException {
+    Throwable cause = e;
+    while (cause != null) {
+      if (cause instanceof DatabricksRetryHandlerException) {
+        throw new DatabricksHttpException(cause.getMessage(), cause);
+      }
+      cause = cause.getCause();
+    }
+    String errorMsg =
+        String.format(
+            "Caught error while executing http request: [%s]. Error Message: [%s]",
+            RequestSanitizer.sanitizeRequest(request), e);
+    LoggingUtil.log(logLevel, errorMsg);
+    throw new DatabricksHttpException(errorMsg, e);
+  }
 }

src/main/java/com/databricks/jdbc/commons/util/DefaultHttpClientUtil.java

@@ -1,7 +1,7 @@
 package com.databricks.jdbc.telemetry;
 
 import com.databricks.jdbc.client.DatabricksHttpException;
-import com.databricks.jdbc.commons.util.DefaultHttpClientUtil;
+import com.databricks.jdbc.client.http.DatabricksHttpClient;
 import com.databricks.jdbc.core.DatabricksSQLException;
 import com.databricks.jdbc.driver.IDatabricksConnectionContext;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -15,7 +15,6 @@
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.utils.URIBuilder;
-import org.apache.http.impl.client.DefaultHttpClient;
 import org.apache.http.util.EntityUtils;
 
 // TODO (Bhuvan) Flush out metrics once the driver connection is closed
@@ -31,7 +30,7 @@ public class DatabricksMetrics {
   private final String METRICS_TYPE = "metrics_type";
   private Boolean hasInitialExportOccurred = false;
   private String workspaceId = null;
-  private DefaultHttpClient telemetryClient = null;
+  private final DatabricksHttpClient telemetryClient;
 
   private void setWorkspaceId(String workspaceId) {
     this.workspaceId = workspaceId;
@@ -68,13 +67,7 @@ public DatabricksMetrics(IDatabricksConnectionContext context) throws Databricks
     }
     String resourceId = context.getComputeResource().getWorkspaceId();
     setWorkspaceId(resourceId);
-    // TODO: Bhuvan: Replace telemetry client with DatabricksHttpClient when non-SSL variation is
-    // implemented
-    try {
-      telemetryClient = DefaultHttpClientUtil.getDefaultHttpClient();
-    } catch (Exception e) {
-      throw new DatabricksSQLException("Failed to create telemetry client");
-    }
+    this.telemetryClient = DatabricksHttpClient.getInstance(context);
     scheduleExportMetrics();
   }
 
@@ -101,7 +94,8 @@ public String sendRequest(Map<String, Double> map, MetricsType metricsType) thro
     uriBuilder.addParameter(METRICS_TYPE, metricsType.name().equals("GAUGE") ? "1" : "0");
     HttpUriRequest request = new HttpPost(uriBuilder.build());
     // TODO (Bhuvan): Add authentication headers
-    CloseableHttpResponse response = telemetryClient.execute(request);
+    // TODO (Bhuvan): execute request using SSL
+    CloseableHttpResponse response = telemetryClient.executeWithoutSSL(request);
 
     // Error handling
     if (response == null) {