Add way to disable refresh tokens

Author: renaudhartert-db

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

@@ -175,6 +175,14 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT")
   private Duration oauthBrowserAuthTimeout;
 
+  /**
+   * Disable automatically adding the offline_access scope to the OAuth authentication
+   * request to request refresh tokens. Note that this does not remove the scope if it 
+   * is explicitly provided by the user.
+   */
+  @ConfigAttribute(env = "DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN")
+  private Boolean disableOauthRefreshToken;
+
   public Environment getEnv() {
     return env;
   }
@@ -631,6 +639,15 @@ public DatabricksConfig setOAuthBrowserAuthTimeout(Duration oauthBrowserAuthTime
     return this;
   }
 
+  public boolean getDisableOauthRefreshToken() {
+    return disableOauthRefreshToken != null && disableOauthRefreshToken;
+  }
+
+  public DatabricksConfig setDisableOauthRefreshToken(boolean disable) {
+    this.disableOauthRefreshToken = disable;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;

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

@@ -9,6 +9,7 @@
 import java.util.HashSet;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.List;
 import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -105,20 +106,25 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
     }
   }
 
-  CachedTokenSource performBrowserAuth(
-      DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache)
-      throws IOException {
-    LOGGER.debug("Performing browser authentication");
-
+  protected List<String> getScopes(DatabricksConfig config) {
     // Get user-provided scopes and add required default scopes.
     Set<String> scopes = new HashSet<>(config.getScopes());
-
-    // Needed to request a refresh token.
-    scopes.add("offline_access");
-
+    // Requesting a refresh token is most of the time the right thing to do to enable
+    // long-lived access to the API. However, some Identity Providers do not support 
+    // refresh tokens.
+    if (!config.getDisableOauthRefreshToken()) {
+      scopes.add("offline_access");
+    }
     if (config.isAzure()) {
       scopes.add(config.getEffectiveAzureLoginAppId() + "/user_impersonation");
     }
+    return new ArrayList<>(scopes);
+  }
+
+  CachedTokenSource performBrowserAuth(
+      DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache)
+      throws IOException {
+    LOGGER.debug("Performing browser authentication");
 
     OAuthClient client =
         new OAuthClient.Builder()
@@ -129,7 +135,7 @@ CachedTokenSource performBrowserAuth(
             .withAccountId(config.getAccountId())
             .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
             .withBrowserTimeout(config.getOAuthBrowserAuthTimeout())
-            .withScopes(new ArrayList<>(scopes))
+            .withScopes(getScopes(config))
             .withOpenIDConnectEndpoints(config.getOidcEndpoints())
             .build();
     Consent consent = client.initiateConsent();

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

@@ -299,4 +299,27 @@ public void testDisableRetriesEnvironmentVariable() {
 
     assertEquals(true, config.getDisableRetries());
   }
+
+  @Test
+  public void testDisableOauthRefreshTokenDefaultValue() {
+    DatabricksConfig config = new DatabricksConfig();
+    assertEquals(false, config.getDisableOauthRefreshToken());
+  }
+
+  @Test
+  public void testDisableOauthRefreshTokenSetAndGet() {
+    DatabricksConfig config = new DatabricksConfig().setDisableOauthRefreshToken(true);
+    assertEquals(true, config.getDisableOauthRefreshToken());
+  }
+
+  @Test
+  public void testDisableOauthRefreshTokenEnvironmentVariable() {
+    Map<String, String> env = new HashMap<>();
+    env.put("DATABRICKS_DISABLE_OAUTH_REFRESH_TOKEN", "true");
+
+    DatabricksConfig config = new DatabricksConfig();
+    config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
+
+    assertEquals(true, config.getDisableOauthRefreshToken());
+  }
 }

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java

@@ -16,6 +16,7 @@
 import java.time.Instant;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import org.junit.jupiter.api.Test;
@@ -551,4 +552,53 @@ void cacheWithInvalidTokensTest() throws IOException {
     // Verify token was saved after browser auth (for the new token)
     Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
   }
+
+  @Test
+  void doNotAddOfflineAccessScopeWhenDisableOauthRefreshTokenIsTrue() {
+    DatabricksConfig config =
+        new DatabricksConfig()
+            .setHost("https://test.databricks.com")
+            .setClientId("test-client-id")
+            .setDisableOauthRefreshToken(true)
+            .setScopes(Arrays.asList("my-test-scope"));
+
+    ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider();
+    List<String> scopes = provider.getScopes(config);
+
+    assertEquals(1, scopes.size());
+    assertTrue(scopes.contains("my-test-scope"));
+  }
+
+  @Test
+  void doNotRemoveUserProvidedScopesWhenDisableOauthRefreshTokenIsTrue() {
+    DatabricksConfig config =
+        new DatabricksConfig()
+            .setHost("https://test.databricks.com")
+            .setClientId("test-client-id")
+            .setDisableOauthRefreshToken(true)
+            .setScopes(Arrays.asList("my-test-scope", "offline_access"));
+
+    ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider();
+    List<String> scopes = provider.getScopes(config);
+
+    assertEquals(2, scopes.size());
+    assertTrue(scopes.contains("offline_access"));
+    assertTrue(scopes.contains("my-test-scope"));
+  }
+
+  @Test
+  void addOfflineAccessScopeWhenDisableOauthRefreshTokenIsFalse() {
+    DatabricksConfig config =
+        new DatabricksConfig()
+            .setHost("https://test.databricks.com")
+            .setClientId("test-client-id")
+            .setScopes(Arrays.asList("my-test-scope"));
+
+    ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider();
+    List<String> scopes = provider.getScopes(config);
+
+    assertEquals(2, scopes.size());
+    assertTrue(scopes.contains("offline_access"));
+    assertTrue(scopes.contains("my-test-scope"));
+  }
 }