Add Keystore support to OSS JDBC Driver

NEXT_CHANGELOG.md

@@ -5,6 +5,8 @@
 ### Added
 - Added support for DoD (.mil) domains
 - Support to fetch metadata in PreparedStatement for SELECT queries before executing the query.
+- Added support for SSL client certificate authentication via keystore configuration parameters: SSLKeyStore, SSLKeyStorePwd, SSLKeyStoreType, and SSLKeyStoreProvider.
+
 
 ### Updated
 - 

src/main/java/com/databricks/jdbc/common/DatabricksClientConfiguratorManager.java

@@ -3,7 +3,7 @@
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.dbclient.impl.common.ClientConfigurator;
 import com.databricks.jdbc.exception.DatabricksDriverException;
-import com.databricks.jdbc.exception.DatabricksHttpException;
+import com.databricks.jdbc.exception.DatabricksSSLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
@@ -28,9 +28,9 @@ public ClientConfigurator getConfigurator(IDatabricksConnectionContext context)
           k -> {
             try {
               return new ClientConfigurator(context);
-            } catch (DatabricksHttpException e) {
+            } catch (DatabricksSSLException e) {
               String message =
-                  String.format("client configurator failed due to HTTP error: %s", e.getMessage());
+                  String.format("client configurator failed due to SSL error: %s", e.getMessage());
               LOGGER.error(e, message);
               throw new DatabricksDriverException(message, DatabricksDriverErrorCode.AUTH_ERROR);
             }

src/main/java/com/databricks/jdbc/dbclient/impl/common/ClientConfigurator.java

@@ -9,8 +9,8 @@
 import com.databricks.jdbc.common.DatabricksJdbcConstants;
 import com.databricks.jdbc.common.util.DatabricksAuthUtil;
 import com.databricks.jdbc.common.util.DriverUtil;
-import com.databricks.jdbc.exception.DatabricksHttpException;
 import com.databricks.jdbc.exception.DatabricksParsingException;
+import com.databricks.jdbc.exception.DatabricksSSLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
@@ -43,7 +43,7 @@ public class ClientConfigurator {
   private DatabricksConfig databricksConfig;
 
   public ClientConfigurator(IDatabricksConnectionContext connectionContext)
-      throws DatabricksHttpException {
+      throws DatabricksSSLException {
     this.connectionContext = connectionContext;
     this.databricksConfig = new DatabricksConfig();
     CommonsHttpClient.Builder httpClientBuilder = new CommonsHttpClient.Builder();
@@ -109,7 +109,7 @@ private static String createUniqueIdentifier(String host, String clientId, List<
    * @param httpClientBuilder The builder to which the SSL configuration should be added.
    */
   void setupConnectionManager(CommonsHttpClient.Builder httpClientBuilder)
-      throws DatabricksHttpException {
+      throws DatabricksSSLException {
     PoolingHttpClientConnectionManager connManager =
         ConfiguratorUtils.getBaseConnectionManager(connectionContext);
     // Default value is 100 which is consistent with the value in the SDK

src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java

@@ -14,6 +14,7 @@
 import com.databricks.jdbc.exception.DatabricksDriverException;
 import com.databricks.jdbc.exception.DatabricksHttpException;
 import com.databricks.jdbc.exception.DatabricksRetryHandlerException;
+import com.databricks.jdbc.exception.DatabricksSSLException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
 import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
@@ -135,7 +136,7 @@ private PoolingHttpClientConnectionManager initializeConnectionManager(
       connectionManager.setMaxTotal(DEFAULT_MAX_HTTP_CONNECTIONS);
       connectionManager.setDefaultMaxPerRoute(DEFAULT_MAX_HTTP_CONNECTIONS_PER_ROUTE);
       return connectionManager;
-    } catch (DatabricksHttpException e) {
+    } catch (DatabricksSSLException e) {
       LOGGER.error("Failed to initialize HTTP connection manager", e);
       // Currently only SSL Handshake failure causes this exception.
       throw new DatabricksDriverException(

src/main/java/com/databricks/jdbc/exception/DatabricksSSLException.java

@@ -0,0 +1,24 @@
+package com.databricks.jdbc.exception;
+
+import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
+
+/** Exception class to handle SSL/TLS configuration and handshake errors. */
+public class DatabricksSSLException extends DatabricksSQLException {
+
+  public DatabricksSSLException(
+      String message, Throwable cause, DatabricksDriverErrorCode sqlCode) {
+    super(message, cause, sqlCode);
+  }
+
+  public DatabricksSSLException(String message, DatabricksDriverErrorCode internalCode) {
+    super(message, null, internalCode.toString());
+  }
+
+  public DatabricksSSLException(String message, String sqlState) {
+    super(message, null, sqlState);
+  }
+
+  public DatabricksSSLException(String message, Throwable throwable, String sqlState) {
+    super(message, throwable, sqlState);
+  }
+}

src/test/java/com/databricks/jdbc/auth/SSLConnectionParametersTest.java

@@ -7,8 +7,8 @@
 import com.databricks.jdbc.api.internal.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.util.SocketFactoryUtil;
 import com.databricks.jdbc.dbclient.impl.common.ConfiguratorUtils;
-import com.databricks.jdbc.exception.DatabricksHttpException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
+import com.databricks.jdbc.exception.DatabricksSSLException;
 import java.security.cert.X509Certificate;
 import java.util.Properties;
 import javax.net.ssl.X509TrustManager;
@@ -32,7 +32,7 @@ public void setUp() {
   }
 
   @Test
-  public void testGetBaseConnectionManagerWithDefaultSettings() throws DatabricksHttpException {
+  public void testGetBaseConnectionManagerWithDefaultSettings() throws DatabricksSSLException {
     when(mockContext.allowSelfSignedCerts()).thenReturn(false);
     when(mockContext.useSystemTrustStore()).thenReturn(false);
     when(mockContext.getSSLTrustStore()).thenReturn(null);
@@ -46,7 +46,7 @@ public void testGetBaseConnectionManagerWithDefaultSettings() throws DatabricksH
   }
 
   @Test
-  public void testGetBaseConnectionManagerWithSelfSignedCerts() throws DatabricksHttpException {
+  public void testGetBaseConnectionManagerWithSelfSignedCerts() throws DatabricksSSLException {
     when(mockContext.allowSelfSignedCerts()).thenReturn(true);
 
     PoolingHttpClientConnectionManager manager =
@@ -68,7 +68,7 @@ public void testGetBaseConnectionManagerWithCustomTrustStore() {
     try {
       ConfiguratorUtils.getBaseConnectionManager(mockContext);
       fail("Should throw exception for non-existent trust store");
-    } catch (DatabricksHttpException e) {
+    } catch (DatabricksSSLException e) {
       assertTrue(
           e.getMessage()
               .contains("Error while setting up custom trust store: /path/to/truststore.jks"),
@@ -88,8 +88,12 @@ public void testGetTrustAllSocketFactoryRegistry() {
 
   @Test
   public void testGetConnectionSocketFactoryRegistryWithSelfSignedCerts()
-      throws DatabricksHttpException {
-    when(mockContext.allowSelfSignedCerts()).thenReturn(true);
+      throws DatabricksSSLException {
+    when(mockContext.allowSelfSignedCerts()).thenReturn(false);
+    when(mockContext.useSystemTrustStore()).thenReturn(false);
+    when(mockContext.getSSLTrustStore()).thenReturn(null);
+    when(mockContext.checkCertificateRevocation()).thenReturn(false);
+    when(mockContext.acceptUndeterminedCertificateRevocation()).thenReturn(false);
 
     Registry<ConnectionSocketFactory> registry =
         ConfiguratorUtils.createConnectionSocketFactoryRegistry(mockContext);

src/test/java/com/databricks/jdbc/dbclient/impl/common/ClientConfiguratorTest.java

@@ -13,9 +13,9 @@
 import com.databricks.jdbc.common.AuthFlow;
 import com.databricks.jdbc.common.AuthMech;
 import com.databricks.jdbc.common.DatabricksJdbcConstants;
-import com.databricks.jdbc.exception.DatabricksHttpException;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.exception.DatabricksSQLException;
+import com.databricks.jdbc.exception.DatabricksSSLException;
 import com.databricks.sdk.WorkspaceClient;
 import com.databricks.sdk.core.DatabricksConfig;
 import com.databricks.sdk.core.DatabricksException;
@@ -40,7 +40,7 @@ public class ClientConfiguratorTest {
 
   @Test
   void getWorkspaceClient_PAT_AuthenticatesWithAccessToken()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.PAT);
     when(mockContext.getHostUrl()).thenReturn("https://pat.databricks.com");
     when(mockContext.getToken()).thenReturn("pat-token");
@@ -58,7 +58,7 @@ void getWorkspaceClient_PAT_AuthenticatesWithAccessToken()
 
   @Test
   void getWorkspaceClient_OAuthWithTokenPassthrough_AuthenticatesCorrectly()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.TOKEN_PASSTHROUGH);
     when(mockContext.getHostUrl()).thenReturn("https://oauth-token.databricks.com");
@@ -77,7 +77,7 @@ void getWorkspaceClient_OAuthWithTokenPassthrough_AuthenticatesCorrectly()
 
   @Test
   void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectly()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS);
     when(mockContext.getHostForOAuth()).thenReturn("https://oauth-client.databricks.com");
@@ -98,7 +98,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectly()
 
   @Test
   void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyGCP()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS);
     when(mockContext.getHostForOAuth()).thenReturn("https://oauth-client.databricks.com");
@@ -118,7 +118,7 @@ void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyGCP()
 
   @Test
   void getWorkspaceClient_OAuthWithClientCredentials_AuthenticatesCorrectlyWithJWT()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getConnectionUuid()).thenReturn("connection-uuid");
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS);
@@ -169,7 +169,7 @@ void testM2MWithJWT() throws DatabricksSQLException {
 
   @Test
   void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrectly()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION);
     when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com");
@@ -194,7 +194,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect
   @Test
   void
       getWorkspaceClient_OAuthWithBrowserBasedAuthentication_WithDiscoveryURL_AuthenticatesCorrectly()
-          throws DatabricksParsingException, IOException, DatabricksHttpException {
+          throws DatabricksParsingException, IOException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION);
     when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com");
@@ -220,7 +220,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_AuthenticatesCorrect
   }
 
   @Test
-  void testNonOauth() throws DatabricksHttpException {
+  void testNonOauth() throws DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OTHER);
     when(mockContext.getHttpConnectionPoolSize()).thenReturn(100);
     configurator = new ClientConfigurator(mockContext);
@@ -250,7 +250,7 @@ void testNonProxyHostsFormatConversion() {
   }
 
   @Test
-  void testSetupProxyConfig() throws DatabricksHttpException {
+  void testSetupProxyConfig() throws DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.PAT);
     when(mockContext.getUseProxy()).thenReturn(true);
     when(mockContext.getProxyHost()).thenReturn("proxy.host.com");
@@ -281,7 +281,7 @@ void testSetupProxyConfig() throws DatabricksHttpException {
 
   @Test
   void setupM2MConfig_WithAzureTenantId_ConfiguresCorrectly()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.CLIENT_CREDENTIALS);
     when(mockContext.getHostForOAuth()).thenReturn("https://azure-oauth.databricks.com");
@@ -445,7 +445,7 @@ void getWorkspaceClient_OAuthWithBrowserBasedAuthentication_SetsCustomRedirectUr
 
   @Test
   void testSetupU2MConfig_WithTokenCache()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION);
     when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com");
@@ -493,7 +493,7 @@ void testSetupU2MConfig_WithTokenCacheNoPassphrase() throws DatabricksParsingExc
 
   @Test
   void testSetupU2MConfig_WithoutTokenCache()
-      throws DatabricksParsingException, DatabricksHttpException {
+      throws DatabricksParsingException, DatabricksSSLException {
     when(mockContext.getAuthMech()).thenReturn(AuthMech.OAUTH);
     when(mockContext.getAuthFlow()).thenReturn(AuthFlow.BROWSER_BASED_AUTHENTICATION);
     when(mockContext.getHostForOAuth()).thenReturn("https://oauth-browser.databricks.com");