Add u2m timeout

Author: renaudhartert-db

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

@@ -1,6 +1,7 @@
 package com.databricks.sdk.core;
 
 import java.lang.reflect.Field;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Objects;
 
@@ -41,16 +42,25 @@ public void setValueOnConfig(DatabricksConfig cfg, String value) throws IllegalA
     // workspace clients or config resolution) are safe
     synchronized (field) {
       field.setAccessible(true);
+      if (value == null || value.trim().isEmpty()) {
+        return;
+      }
+      
       if (field.getType() == String.class) {
         field.set(cfg, value);
       } else if (field.getType() == int.class) {
         field.set(cfg, Integer.parseInt(value));
+      } else if (field.getType() == Integer.class) {
+        field.set(cfg, Integer.parseInt(value));
       } else if (field.getType() == boolean.class) {
         field.set(cfg, Boolean.parseBoolean(value));
+      } else if (field.getType() == Boolean.class) {
+        field.set(cfg, Boolean.parseBoolean(value));
+      } else if (field.getType() == Duration.class) {
+        int seconds = Integer.parseInt(value);
+        field.set(cfg, seconds > 0 ? Duration.ofSeconds(seconds) : null);
       } else if (field.getType() == ProxyConfig.ProxyAuthType.class) {
-        if (value != null) {
-          field.set(cfg, ProxyConfig.ProxyAuthType.valueOf(value));
-        }
+        field.set(cfg, ProxyConfig.ProxyAuthType.valueOf(value));
       }
       field.setAccessible(false);
     }

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

@@ -14,6 +14,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.time.Duration;
 import java.util.*;
 import org.apache.http.HttpMessage;
 
@@ -163,6 +164,13 @@ public class DatabricksConfig {
   @ConfigAttribute(env = "DATABRICKS_DISABLE_ASYNC_TOKEN_REFRESH")
   private Boolean disableAsyncTokenRefresh;
 
+  /**
+   * The duration to wait for a browser response during U2M authentication before timing out. 
+   * If set to 0 or null, the connector waits for an indefinite amount of time.
+   */
+  @ConfigAttribute(env = "DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT")
+  private Duration oauthBrowserAuthTimeout;
+
   public Environment getEnv() {
     return env;
   }
@@ -588,6 +596,20 @@ public DatabricksConfig setDisableAsyncTokenRefresh(boolean disableAsyncTokenRef
     return this;
   }
 
+  public Duration getOAuthBrowserAuthTimeout() {
+    return oauthBrowserAuthTimeout;
+  }
+
+  public DatabricksConfig setOAuthBrowserAuthTimeout(Duration oauthBrowserAuthTimeout) {
+    this.oauthBrowserAuthTimeout = oauthBrowserAuthTimeout;
+    return this;
+  }
+
+  public DatabricksConfig setOAuthBrowserAuthTimeout(int seconds) {
+    this.oauthBrowserAuthTimeout = seconds > 0 ? Duration.ofSeconds(seconds) : null;
+    return this;
+  }
+
   public boolean isAzure() {
     if (azureWorkspaceResourceId != null) {
       return true;

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

@@ -18,6 +18,9 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
+import java.time.Duration;
+
 import org.apache.commons.io.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,6 +56,7 @@ public class Consent implements Serializable {
   private final String redirectUrl;
   private final String clientId;
   private final String clientSecret;
+  private final Optional<Duration> browserTimeout;
 
   public static class Builder {
     private HttpClient hc = new CommonsHttpClient.Builder().withTimeoutSeconds(30).build();
@@ -63,6 +67,7 @@ public static class Builder {
     private String redirectUrl;
     private String clientId;
     private String clientSecret;
+    private Optional<Duration> browserTimeout = Optional.empty();
 
     public Builder withHttpClient(HttpClient hc) {
       this.hc = hc;
@@ -104,6 +109,15 @@ public Builder withClientSecret(String clientSecret) {
       return this;
     }
 
+    public Builder withBrowserTimeout(Optional<Duration> browserTimeout) {
+      this.browserTimeout = browserTimeout;
+      return this;
+    }
+
+    public Builder withBrowserTimeout(Duration browserTimeout) {
+      return withBrowserTimeout(Optional.of(browserTimeout));
+    }
+
     public Consent build() {
       return new Consent(this);
     }
@@ -119,6 +133,7 @@ private Consent(Builder builder) {
     this.clientId = Objects.requireNonNull(builder.clientId);
     // This may be null for native apps or single-page apps.
     this.clientSecret = builder.clientSecret;
+    this.browserTimeout = builder.browserTimeout;
   }
 
   public Consent setHttpClient(HttpClient hc) {
@@ -155,6 +170,10 @@ public String getClientSecret() {
     return clientSecret;
   }
 
+  public Optional<Duration> getBrowserTimeout() {
+    return browserTimeout;
+  }
+
   /**
    * Launch a browser to collect an authorization code and exchange the code for an OAuth token.
    *
@@ -219,7 +238,8 @@ private Map<String, String> getOAuthCallbackParameters() throws IOException {
               + redirect.getHost()
               + ", redirectUrl host must be one of: localhost, 127.0.0.1");
     }
-    CallbackResponseHandler handler = new CallbackResponseHandler();
+
+    CallbackResponseHandler handler = new CallbackResponseHandler(this.browserTimeout);
     HttpServer httpServer =
         HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0);
     httpServer.createContext("/", handler);
@@ -285,6 +305,11 @@ static class CallbackResponseHandler implements HttpHandler {
     // Protects params
     private final Object lock = new Object();
     private volatile Map<String, String> params;
+    private final Optional<Duration> timeout;
+
+    public CallbackResponseHandler(Optional<Duration> timeout) {
+      this.timeout = timeout;
+    }
 
     @Override
     public void handle(HttpExchange exchange) {
@@ -373,7 +398,15 @@ public Map<String, String> getParams() {
       synchronized (lock) {
         if (params == null) {
           try {
-            lock.wait();
+            if (timeout.isPresent()) {
+              Duration t = timeout.get();
+              lock.wait(t.toMillis());
+              if (params == null) {
+                throw new DatabricksException("OAuth browser authentication timed out after " + t.getSeconds() + " seconds");
+              }
+            } else {
+              lock.wait();
+            }
           } catch (InterruptedException e) {
             throw new DatabricksException(
                 "Interrupted while waiting for parameters: " + e.getMessage(), e);

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

@@ -7,6 +7,7 @@
 import java.nio.file.Path;
 import java.util.Objects;
 import java.util.Optional;
+import java.time.Duration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -115,6 +116,7 @@ CachedTokenSource performBrowserAuth(
             .withAccountId(config.getAccountId())
             .withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
             .withScopes(config.getScopes())
+            .withBrowserTimeout(config.getOAuthBrowserAuthTimeout())
             .build();
     Consent consent = client.initiateConsent();
 

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

@@ -9,7 +9,14 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.security.SecureRandom;
-import java.util.*;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Base64;
+import java.util.HashMap;
 import java.util.stream.Collectors;
 
 /**
@@ -36,6 +43,7 @@ public static class Builder {
     private String clientSecret;
     private HttpClient hc;
     private String accountId;
+    private Optional<Duration> browserTimeout = Optional.empty();
 
     public Builder() {}
 
@@ -77,6 +85,11 @@ public Builder withAccountId(String accountId) {
       this.accountId = accountId;
       return this;
     }
+
+    public Builder withBrowserTimeout(Duration browserTimeout) {
+      this.browserTimeout = Optional.of(browserTimeout);
+      return this;
+    }
   }
 
   private final String clientId;
@@ -90,6 +103,7 @@ public Builder withAccountId(String accountId) {
   private final SecureRandom random = new SecureRandom();
   private final boolean isAws;
   private final boolean isAzure;
+  private final Optional<Duration> browserTimeout;
 
   private OAuthClient(Builder b) throws IOException {
     this.clientId = Objects.requireNonNull(b.clientId);
@@ -120,6 +134,7 @@ private OAuthClient(Builder b) throws IOException {
               config.getEffectiveAzureLoginAppId() + "/user_impersonation", "offline_access");
     }
     this.scopes = scopes;
+    this.browserTimeout = b.browserTimeout;
   }
 
   public String getHost() {
@@ -207,6 +222,7 @@ public Consent initiateConsent() throws MalformedURLException {
         .withState(state)
         .withVerifier(verifier)
         .withHttpClient(hc)
+        .withBrowserTimeout(browserTimeout.orElse(Duration.ofSeconds(0)))
         .build();
   }
 }

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

@@ -12,6 +12,7 @@
 import com.databricks.sdk.core.oauth.TokenSource;
 import com.databricks.sdk.core.utils.Environment;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -250,4 +251,62 @@ public void testGetTokenSourceWithOAuth() {
     assertFalse(tokenSource instanceof ErrorTokenSource);
     assertEquals(tokenSource.getToken().getAccessToken(), "test-token");
   }
+
+  @Test
+  public void testOAuthBrowserAuthTimeout() {
+    DatabricksConfig config = new DatabricksConfig();
+    
+    // Test default value (should be null)
+    assertNull(config.getOAuthBrowserAuthTimeout());
+    
+    // Test setting the timeout using Duration
+    config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(30));
+    assertEquals(Duration.ofSeconds(30), config.getOAuthBrowserAuthTimeout());
+    
+    // Test setting the timeout using int (backward compatibility)
+    config.setOAuthBrowserAuthTimeout(60);
+    assertEquals(Duration.ofSeconds(60), config.getOAuthBrowserAuthTimeout());
+    
+    // Test setting to 0 (should be null for indefinite wait)
+    config.setOAuthBrowserAuthTimeout(0);
+    assertNull(config.getOAuthBrowserAuthTimeout());
+  }
+
+  @Test
+  public void testEnvironmentVariableLoading() {
+    // Test that environment variables are now loaded properly for Duration and Integer
+    DatabricksConfig config = new DatabricksConfig();
+    
+    // Set environment variable
+    Map<String, String> env = new HashMap<>();
+    env.put("DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT", "30");
+    env.put("DATABRICKS_DEBUG_TRUNCATE_BYTES", "100");
+    env.put("DATABRICKS_RATE_LIMIT", "50");
+    
+    // Resolve config with environment
+    Environment environment = new Environment(env, new ArrayList<>(), "Linux");
+    config.resolve(environment);
+    
+    // These should now be loaded properly
+    assertEquals(Duration.ofSeconds(30), config.getOAuthBrowserAuthTimeout());
+    assertEquals(Integer.valueOf(100), config.getDebugTruncateBytes());
+    assertEquals(Integer.valueOf(50), config.getRateLimit());
+  }
+
+  @Test
+  public void testEnvironmentVariableLoadingZeroTimeout() {
+    // Test that environment variable with value 0 results in null Duration
+    DatabricksConfig config = new DatabricksConfig();
+    
+    // Set environment variable to 0
+    Map<String, String> env = new HashMap<>();
+    env.put("DATABRICKS_OAUTH_BROWSER_AUTH_TIMEOUT", "0");
+    
+    // Resolve config with environment
+    Environment environment = new Environment(env, new ArrayList<>(), "Linux");
+    config.resolve(environment);
+    
+    // Should be null for indefinite wait
+    assertNull(config.getOAuthBrowserAuthTimeout());
+  }
 }

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

@@ -0,0 +1,43 @@
+package com.databricks.sdk.core.oauth;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.time.Duration;
+import java.util.Optional;
+import org.junit.jupiter.api.Test;
+
+public class ConsentTest {
+
+  @Test
+  public void testConsentWithBrowserAuthTimeout() {
+    Consent consent = new Consent.Builder()
+        .withClientId("test-client-id")
+        .withClientSecret("test-client-secret")
+        .withAuthUrl("https://test.com/auth")
+        .withTokenUrl("https://test.com/token")
+        .withRedirectUrl("http://localhost:8080/callback")
+        .withState("test-state")
+        .withVerifier("test-verifier")
+        .withBrowserTimeout(Duration.ofSeconds(30))
+        .build();
+
+    // Verify that the timeout is properly set
+    assertEquals(Optional.of(Duration.ofSeconds(30)), consent.getBrowserTimeout());
+  }
+
+  @Test
+  public void testConsentWithoutBrowserAuthTimeout() {
+    Consent consent = new Consent.Builder()
+        .withClientId("test-client-id")
+        .withClientSecret("test-client-secret")
+        .withAuthUrl("https://test.com/auth")
+        .withTokenUrl("https://test.com/token")
+        .withRedirectUrl("http://localhost:8080/callback")
+        .withState("test-state")
+        .withVerifier("test-verifier")
+        .build();
+
+    // Verify that the timeout is empty when not set
+    assertEquals(Optional.empty(), consent.getBrowserTimeout());
+  }
+}