[Draft] Add cache for oidc endpoint

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -12,6 +12,7 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import org.apache.http.HttpMessage;
 
 public class DatabricksConfig {
@@ -139,6 +140,7 @@ public class DatabricksConfig {
 
   private Environment env;
 
+  ConcurrentHashMap<String, OpenIDConnectEndpoints> oidcEndpointsCache = new ConcurrentHashMap<>();
   private DatabricksEnvironment databricksEnvironment;
 
   public Environment getEnv() {
@@ -374,13 +376,17 @@ public DatabricksConfig setAzureUseMsi(boolean azureUseMsi) {
     return this;
   }
 
-  /** @deprecated Use {@link #getAzureUseMsi()} instead. */
+  /**
+   * @deprecated Use {@link #getAzureUseMsi()} instead.
+   */
   @Deprecated()
   public boolean getAzureUseMSI() {
     return azureUseMsi;
   }
 
-  /** @deprecated Use {@link #setAzureUseMsi(boolean)} instead. */
+  /**
+   * @deprecated Use {@link #setAzureUseMsi(boolean)} instead.
+   */
   @Deprecated
   public DatabricksConfig setAzureUseMSI(boolean azureUseMsi) {
     this.azureUseMsi = azureUseMsi;
@@ -545,20 +551,7 @@ public OpenIDConnectEndpoints getOidcEndpoints() throws IOException {
     if (discoveryUrl == null) {
       return fetchDefaultOidcEndpoints();
     }
-    return fetchOidcEndpointsFromDiscovery();
-  }
-
-  private OpenIDConnectEndpoints fetchOidcEndpointsFromDiscovery() {
-    try {
-      Request request = new Request("GET", discoveryUrl);
-      Response resp = getHttpClient().execute(request);
-      if (resp.getStatusCode() == 200) {
-        return new ObjectMapper().readValue(resp.getBody(), OpenIDConnectEndpoints.class);
-      }
-    } catch (IOException e) {
-      throw ConfigLoader.makeNicerError(e.getMessage(), e, this);
-    }
-    return null;
+    return fetchOidcEndpointWithCache(discoveryUrl);
   }
 
   private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException {
@@ -582,11 +575,7 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException {
     }
 
     String oidcEndpoint = getHost() + "/oidc/.well-known/oauth-authorization-server";
-    Response resp = getHttpClient().execute(new Request("GET", oidcEndpoint));
-    if (resp.getStatusCode() != 200) {
-      return null;
-    }
-    return new ObjectMapper().readValue(resp.getBody(), OpenIDConnectEndpoints.class);
+    return fetchOidcEndpointWithCache(oidcEndpoint);
   }
 
   @Override
@@ -656,4 +645,20 @@ public DatabricksConfig newWithWorkspaceHost(String host) {
     newConfig.setHost(host);
     return newConfig;
   }
+
+  private OpenIDConnectEndpoints fetchOidcEndpointWithCache(String oidcEndpoint) {
+    return oidcEndpointsCache.computeIfAbsent(
+        oidcEndpoint,
+        key -> {
+          try {
+            Response resp = getHttpClient().execute(new Request("GET", oidcEndpoint));
+            if (resp.getStatusCode() != 200) {
+              return null;
+            }
+            return new ObjectMapper().readValue(resp.getBody(), OpenIDConnectEndpoints.class);
+          } catch (IOException e) {
+            throw ConfigLoader.makeNicerError(e.getMessage(), e, this);
+          }
+        });
+  }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java

@@ -26,7 +26,6 @@ public HeaderFactory configure(DatabricksConfig config) {
         || config.getHost() == null) {
       return null;
     }
-    // TODO: Add cache in getOidcEndpoints
     // TODO: Azure returns 404 for UC workspace after redirecting to
     // https://login.microsoftonline.com/{cfg.azure_tenant_id}/.well-known/oauth-authorization-server
     try {