Add DiscoveryURL to flows other than M2M (#599)

* Add OAuth to flows other than M2M

* Add tests

Author: samikshya-db

src/main/java/com/databricks/jdbc/auth/OAuthEndpointResolver.java

@@ -5,6 +5,7 @@
 
 import com.databricks.jdbc.api.IDatabricksConnectionContext;
 import com.databricks.jdbc.common.DatabricksJdbcConstants;
+import com.databricks.jdbc.common.util.DatabricksAuthUtil;
 import com.databricks.jdbc.exception.DatabricksParsingException;
 import com.databricks.jdbc.log.JdbcLogger;
 import com.databricks.jdbc.log.JdbcLoggerFactory;
@@ -16,7 +17,6 @@
 import com.databricks.sdk.core.oauth.AuthParameterPosition;
 import com.databricks.sdk.core.oauth.RefreshableTokenSource;
 import com.databricks.sdk.core.oauth.Token;
-import com.google.common.annotations.VisibleForTesting;
 import java.time.LocalDateTime;
 import java.util.HashMap;
 import java.util.Map;
@@ -32,10 +32,9 @@ public class OAuthRefreshCredentialsProvider extends RefreshableTokenSource
   private final String clientId;
   private final String clientSecret;
 
-  @VisibleForTesting
   public OAuthRefreshCredentialsProvider(
-      IDatabricksConnectionContext context, OAuthEndpointResolver oAuthEndpointResolver) {
-    this.tokenEndpoint = oAuthEndpointResolver.getTokenEndpoint();
+      IDatabricksConnectionContext context, DatabricksConfig databricksConfig) {
+    this.tokenEndpoint = DatabricksAuthUtil.getTokenEndpoint(databricksConfig, context);
     try {
       this.clientId = context.getClientId();
     } catch (DatabricksParsingException e) {
@@ -53,11 +52,6 @@ public OAuthRefreshCredentialsProvider(
             LocalDateTime.now().minusMinutes(1));
   }
 
-  public OAuthRefreshCredentialsProvider(
-      IDatabricksConnectionContext context, DatabricksConfig databricksConfig) {
-    this(context, new OAuthEndpointResolver(context, databricksConfig));
-  }
-
   @Override
   public String authType() {
     return "oauth-refresh";

src/main/java/com/databricks/jdbc/auth/OAuthRefreshCredentialsProvider.java

@@ -1,6 +1,7 @@
 package com.databricks.jdbc.auth;
 
 import com.databricks.jdbc.api.IDatabricksConnectionContext;
+import com.databricks.jdbc.common.util.DatabricksAuthUtil;
 import com.databricks.jdbc.dbclient.IDatabricksHttpClient;
 import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory;
 import com.databricks.sdk.core.CredentialsProvider;
@@ -22,15 +23,9 @@ public class PrivateKeyClientCredentialProvider implements CredentialsProvider {
 
   public PrivateKeyClientCredentialProvider(
       IDatabricksConnectionContext connectionContext, DatabricksConfig databricksConfig) {
-    this(connectionContext, new OAuthEndpointResolver(connectionContext, databricksConfig));
-  }
-
-  @VisibleForTesting
-  public PrivateKeyClientCredentialProvider(
-      IDatabricksConnectionContext connectionContext, OAuthEndpointResolver oAuthEndpointResolver) {
     this.connectionContext = connectionContext;
     this.httpClient = DatabricksHttpClientFactory.getInstance().getClient(connectionContext);
-    this.tokenEndpoint = oAuthEndpointResolver.getTokenEndpoint();
+    this.tokenEndpoint = DatabricksAuthUtil.getTokenEndpoint(databricksConfig, connectionContext);
   }
 
   @Override

src/main/java/com/databricks/jdbc/auth/PrivateKeyClientCredentialProvider.java

@@ -0,0 +1,29 @@
+package com.databricks.jdbc.common.util;
+
+import com.databricks.jdbc.api.IDatabricksConnectionContext;
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.databricks.jdbc.model.client.thrift.generated.*;
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.DatabricksException;
+import java.io.IOException;
+import java.util.*;
+
+public class DatabricksAuthUtil {
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(DatabricksAuthUtil.class);
+
+  public static String getTokenEndpoint(
+      DatabricksConfig databricksConfig, IDatabricksConnectionContext connectionContext) {
+    String userProvidedTokenEndpoint = connectionContext.getTokenEndpoint();
+    if (userProvidedTokenEndpoint != null) {
+      return userProvidedTokenEndpoint;
+    }
+    try {
+      return databricksConfig.getOidcEndpoints().getTokenEndpoint();
+    } catch (IOException e) {
+      String errorMessage = "Failed to build default token endpoint URL.";
+      LOGGER.error(errorMessage);
+      throw new DatabricksException(errorMessage, e);
+    }
+  }
+}

src/main/java/com/databricks/jdbc/common/util/DatabricksAuthUtil.java

@@ -38,6 +38,7 @@ public ClientConfigurator(IDatabricksConnectionContext connectionContext) {
     setupProxyConfig(httpClientBuilder);
     setupConnectionManager(httpClientBuilder);
     this.databricksConfig.setHttpClient(httpClientBuilder.build());
+    setupDiscoveryEndpoint();
     setupAuthConfig();
     this.databricksConfig.resolve();
   }
@@ -219,4 +220,10 @@ public static String convertNonProxyHostConfigToBeSystemPropertyCompliant(String
   public DatabricksConfig getDatabricksConfig() {
     return this.databricksConfig;
   }
+
+  private void setupDiscoveryEndpoint() {
+    if (connectionContext.isOAuthDiscoveryModeEnabled()) {
+      databricksConfig.setDiscoveryUrl(connectionContext.getOAuthDiscoveryURL());
+    }
+  }
 }

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

@@ -5,6 +5,9 @@
 import com.databricks.jdbc.common.IDatabricksComputeResource;
 import com.databricks.jdbc.common.Warehouse;
 import com.databricks.jdbc.model.client.thrift.generated.*;
+import com.databricks.sdk.core.DatabricksException;
+import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import java.net.MalformedURLException;
 import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.List;
@@ -102,11 +105,18 @@ public class TestConstants {
   public static final String TEST_JWT_ALGORITHM = "RS256";
   public static final String TEST_JWT_KEY_FILE = "src/test/resources/private_key.pem";
   public static final String TEST_ACCESS_TOKEN = "test-access-token";
-  public static final String TEST_OIDC_RESPONSE =
-      "{\n"
-          + "  \"token_endpoint\": \"https://test.token.url\",\n"
-          + "  \"authorization_endpoint\": \"https://test.auth.url\"\n"
-          + "}";
+
+  public static OpenIDConnectEndpoints TEST_OIDC_ENDPOINTS;
+
+  static {
+    try {
+      TEST_OIDC_ENDPOINTS =
+          new OpenIDConnectEndpoints("https://test.token.url", "https://test.auth.url");
+    } catch (MalformedURLException e) {
+      throw new DatabricksException("Can't initiate test constant for OIDC. Error: " + e);
+    }
+  }
+
   public static final String TEST_OAUTH_RESPONSE =
       "{\n"
           + "  \"expires_in\": 3600,\n"