From 266bc8edac71f1ad187d6bdaaaef383c8ccd866d Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 7 Aug 2025 14:15:46 +0000 Subject: [PATCH 1/9] Add u2m timeout --- .../sdk/core/ConfigAttributeAccessor.java | 16 ++++- .../databricks/sdk/core/DatabricksConfig.java | 22 +++++++ .../databricks/sdk/core/oauth/Consent.java | 37 +++++++++++- .../ExternalBrowserCredentialsProvider.java | 2 + .../sdk/core/oauth/OAuthClient.java | 18 +++++- .../sdk/core/DatabricksConfigTest.java | 59 +++++++++++++++++++ .../sdk/core/oauth/ConsentTest.java | 43 ++++++++++++++ 7 files changed, 191 insertions(+), 6 deletions(-) create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java index 73cb3cba2..1abfece5d 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java +++ b/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); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 0bc0b868a..174602fad 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/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; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index aee9fe50f..4282c82c2 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/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 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 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 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 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 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 params; + private final Optional timeout; + + public CallbackResponseHandler(Optional timeout) { + this.timeout = timeout; + } @Override public void handle(HttpExchange exchange) { @@ -373,7 +398,15 @@ public Map 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); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java index 70b7e1e21..e8f0c3a2b 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java +++ b/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(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java index 139031c8c..f64453c62 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java +++ b/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 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 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(); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index b3ac333a3..600f8a54e 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/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 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 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()); + } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java new file mode 100644 index 000000000..ed9202696 --- /dev/null +++ b/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()); + } +} \ No newline at end of file From 3b06457f346438e61f9d9c8b17a329dc7d92ab56 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 7 Aug 2025 14:18:33 +0000 Subject: [PATCH 2/9] Add changelog + fmt --- NEXT_CHANGELOG.md | 2 + .../sdk/core/ConfigAttributeAccessor.java | 2 +- .../databricks/sdk/core/DatabricksConfig.java | 4 +- .../databricks/sdk/core/oauth/Consent.java | 6 +-- .../ExternalBrowserCredentialsProvider.java | 1 - .../sdk/core/oauth/OAuthClient.java | 5 +-- .../sdk/core/DatabricksConfigTest.java | 20 ++++----- .../sdk/core/oauth/ConsentTest.java | 42 ++++++++++--------- 8 files changed, 42 insertions(+), 40 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index b60768575..20502a976 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -4,6 +4,8 @@ ### New Features and Improvements +* Add option to add a timeout for browser confirmation in the U2M authentication flow. + ### Bug Fixes * User provided scopes are now properly propagated in OAuth flows. diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java index 1abfece5d..06502b508 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java @@ -45,7 +45,7 @@ public void setValueOnConfig(DatabricksConfig cfg, String value) throws IllegalA if (value == null || value.trim().isEmpty()) { return; } - + if (field.getType() == String.class) { field.set(cfg, value); } else if (field.getType() == int.class) { diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 637f2591d..834d58594 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -165,8 +165,8 @@ public class DatabricksConfig { 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. + * 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; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 4282c82c2..e7972f5f3 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -14,13 +14,12 @@ import java.io.Serializable; import java.net.*; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.Arrays; 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; @@ -402,7 +401,8 @@ public Map getParams() { Duration t = timeout.get(); lock.wait(t.toMillis()); if (params == null) { - throw new DatabricksException("OAuth browser authentication timed out after " + t.getSeconds() + " seconds"); + throw new DatabricksException( + "OAuth browser authentication timed out after " + t.getSeconds() + " seconds"); } } else { lock.wait(); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java index f755597c0..95780b6fa 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java @@ -9,7 +9,6 @@ import java.util.HashSet; import java.util.Objects; import java.util.Optional; -import java.time.Duration; import java.util.Set; import org.slf4j.Logger; import org.slf4j.LoggerFactory; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java index 4b395b51b..dc01020b6 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java @@ -10,13 +10,12 @@ import java.security.NoSuchAlgorithmException; import java.security.SecureRandom; import java.time.Duration; -import java.util.Arrays; +import java.util.Base64; +import java.util.HashMap; 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; /** diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 600f8a54e..2d3b7ede3 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -255,18 +255,18 @@ public void testGetTokenSourceWithOAuth() { @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()); @@ -276,17 +276,17 @@ public void testOAuthBrowserAuthTimeout() { public void testEnvironmentVariableLoading() { // Test that environment variables are now loaded properly for Duration and Integer DatabricksConfig config = new DatabricksConfig(); - + // Set environment variable Map 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()); @@ -297,15 +297,15 @@ public void testEnvironmentVariableLoading() { public void testEnvironmentVariableLoadingZeroTimeout() { // Test that environment variable with value 0 results in null Duration DatabricksConfig config = new DatabricksConfig(); - + // Set environment variable to 0 Map 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()); } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java index ed9202696..1ea7a40d0 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java @@ -10,16 +10,17 @@ 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(); + 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()); @@ -27,17 +28,18 @@ public void testConsentWithBrowserAuthTimeout() { @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(); + 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()); } -} \ No newline at end of file +} From 414686b3c2a554db5faaa54994f4fc64812fda8b Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 7 Aug 2025 14:21:16 +0000 Subject: [PATCH 3/9] Simplify --- .../src/main/java/com/databricks/sdk/core/oauth/Consent.java | 4 ---- .../main/java/com/databricks/sdk/core/oauth/OAuthClient.java | 2 +- .../test/java/com/databricks/sdk/core/oauth/ConsentTest.java | 2 +- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index e7972f5f3..f6f7ab6e1 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -113,10 +113,6 @@ public Builder withBrowserTimeout(Optional browserTimeout) { return this; } - public Builder withBrowserTimeout(Duration browserTimeout) { - return withBrowserTimeout(Optional.of(browserTimeout)); - } - public Consent build() { return new Consent(this); } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java index dc01020b6..3803bafb0 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClient.java @@ -211,7 +211,7 @@ public Consent initiateConsent() throws MalformedURLException { .withState(state) .withVerifier(verifier) .withHttpClient(hc) - .withBrowserTimeout(browserTimeout.orElse(Duration.ofSeconds(0))) + .withBrowserTimeout(browserTimeout) .build(); } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java index 1ea7a40d0..f5082cf0c 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java @@ -19,7 +19,7 @@ public void testConsentWithBrowserAuthTimeout() { .withRedirectUrl("http://localhost:8080/callback") .withState("test-state") .withVerifier("test-verifier") - .withBrowserTimeout(Duration.ofSeconds(30)) + .withBrowserTimeout(Optional.of(Duration.ofSeconds(30))) .build(); // Verify that the timeout is properly set From c5057f249d7095f7e96c7154fad90bc3842824c6 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 7 Aug 2025 14:29:09 +0000 Subject: [PATCH 4/9] Adjust tests --- .../sdk/core/DatabricksConfigTest.java | 31 ++----------------- 1 file changed, 2 insertions(+), 29 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index 2d3b7ede3..e672ba1f0 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -256,57 +256,30 @@ public void testGetTokenSourceWithOAuth() { 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 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); + DatabricksConfig config = new DatabricksConfig(); + config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name"))); - // 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 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()); - } } From a0217fbf136f42dde5160218e06a8d139e687304 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 7 Aug 2025 14:48:45 +0000 Subject: [PATCH 5/9] Add tests --- .../databricks/sdk/core/oauth/Consent.java | 15 +-- .../sdk/core/oauth/ConsentTest.java | 91 ++++++++++++++++++- 2 files changed, 98 insertions(+), 8 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index f6f7ab6e1..4c4fa32ad 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -297,8 +297,7 @@ private Token exchange(String code, String state) { static class CallbackResponseHandler implements HttpHandler { private final Logger LOG = LoggerFactory.getLogger(getClass().getName()); - // Protects params - private final Object lock = new Object(); + protected final Object lock = new Object(); // protected for testing private volatile Map params; private final Optional timeout; @@ -343,10 +342,7 @@ public void handleInner(HttpExchange exchange) throws IOException { }); sendSuccess(exchange); - synchronized (lock) { - params = theseParams; - lock.notify(); - } + setParams(theseParams); } private void sendError( @@ -411,5 +407,12 @@ public Map getParams() { return params; } } + + void setParams(Map params) { + synchronized (lock) { + this.params = params; + lock.notify(); + } + } } } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java index f5082cf0c..8ea7a8d3d 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java @@ -2,8 +2,13 @@ import static org.junit.jupiter.api.Assertions.*; +import com.databricks.sdk.core.DatabricksException; import java.time.Duration; +import java.util.HashMap; +import java.util.Map; import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.Test; public class ConsentTest { @@ -22,7 +27,6 @@ public void testConsentWithBrowserAuthTimeout() { .withBrowserTimeout(Optional.of(Duration.ofSeconds(30))) .build(); - // Verify that the timeout is properly set assertEquals(Optional.of(Duration.ofSeconds(30)), consent.getBrowserTimeout()); } @@ -39,7 +43,90 @@ public void testConsentWithoutBrowserAuthTimeout() { .withVerifier("test-verifier") .build(); - // Verify that the timeout is empty when not set assertEquals(Optional.empty(), consent.getBrowserTimeout()); } + + @Test + public void testTimeoutLogicWithShortTimeout() throws InterruptedException { + // Test that timeout is enforced correctly. + Consent.CallbackResponseHandler handler = + new Consent.CallbackResponseHandler(Optional.of(Duration.ofMillis(100))); // 100ms timeout + + long startTime = System.currentTimeMillis(); + + try { + handler.getParams(); + fail("Expected timeout exception"); + } catch (DatabricksException e) { + long elapsedTime = System.currentTimeMillis() - startTime; + assertTrue( + elapsedTime >= 100, "Timeout should have taken at least 100ms, but took " + elapsedTime); + assertTrue(e.getMessage().contains("timed out after 0 seconds")); + } + } + + @Test + public void testTimeoutLogicWithNoTimeout() throws InterruptedException { + // Test that no timeout means indefinite wait. + Consent.CallbackResponseHandler handler = new Consent.CallbackResponseHandler(Optional.empty()); + + CountDownLatch latch = new CountDownLatch(1); + + Thread setterThread = + new Thread( + () -> { + try { + Thread.sleep(50); + synchronized (handler.lock) { + Map params = new HashMap<>(); + params.put("code", "test-code"); + params.put("state", "test-state"); + handler.setParams(params); + } + latch.countDown(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + + setterThread.start(); + + Map result = handler.getParams(); + assertNotNull(result); + assertEquals("test-code", result.get("code")); + assertEquals("test-state", result.get("state")); + assertTrue(latch.await(1, TimeUnit.SECONDS)); + } + + @Test + public void testTimeoutLogicWithParamsSetBeforeTimeout() throws InterruptedException { + // Test that if params are set before timeout, no exception is thrown. + Consent.CallbackResponseHandler handler = + new Consent.CallbackResponseHandler(Optional.of(Duration.ofSeconds(1))); + + CountDownLatch latch = new CountDownLatch(1); + + Thread setterThread = + new Thread( + () -> { + try { + Thread.sleep(50); + synchronized (handler.lock) { + Map params = new HashMap<>(); + params.put("code", "test-code"); + handler.setParams(params); + } + latch.countDown(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + + setterThread.start(); + + Map result = handler.getParams(); + assertNotNull(result); + assertEquals("test-code", result.get("code")); + assertTrue(latch.await(1, TimeUnit.SECONDS)); + } } From 346ffae7d7c834fe40ff9456017ebe7072e576cf Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 7 Aug 2025 14:56:59 +0000 Subject: [PATCH 6/9] Remove duplicate --- .../main/java/com/databricks/sdk/core/DatabricksConfig.java | 5 ----- .../java/com/databricks/sdk/core/DatabricksConfigTest.java | 6 +++--- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java index 834d58594..074e97974 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java @@ -614,11 +614,6 @@ public DatabricksConfig setOAuthBrowserAuthTimeout(Duration oauthBrowserAuthTime 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; diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java index e672ba1f0..88a466a32 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java @@ -261,11 +261,11 @@ public void testOAuthBrowserAuthTimeout() { config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(30)); assertEquals(Duration.ofSeconds(30), config.getOAuthBrowserAuthTimeout()); - config.setOAuthBrowserAuthTimeout(60); + config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(60)); assertEquals(Duration.ofSeconds(60), config.getOAuthBrowserAuthTimeout()); - config.setOAuthBrowserAuthTimeout(0); - assertNull(config.getOAuthBrowserAuthTimeout()); + config.setOAuthBrowserAuthTimeout(Duration.ofSeconds(0)); + assertEquals(Duration.ZERO, config.getOAuthBrowserAuthTimeout()); } @Test From f751d11999d67fbd368bcec586f6efc21c432cff Mon Sep 17 00:00:00 2001 From: Renaud Hartert Date: Thu, 7 Aug 2025 19:44:32 +0200 Subject: [PATCH 7/9] Update Consent.java --- .../java/com/databricks/sdk/core/oauth/Consent.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 4c4fa32ad..73df7cd29 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -239,10 +239,13 @@ private Map getOAuthCallbackParameters() throws IOException { HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0); httpServer.createContext("/", handler); httpServer.start(); - desktopBrowser(); - Map params = handler.getParams(); - httpServer.stop(0); - return params; + + try { + desktopBrowser(); + return handler.getParams(); + } finally { + httpServer.stop(0); + } } /** From 2e2bc6559e2d78f5ece66eb800f88ce40e74a345 Mon Sep 17 00:00:00 2001 From: Renaud Hartert Date: Thu, 7 Aug 2025 20:37:56 +0200 Subject: [PATCH 8/9] Update Consent.java --- .../src/main/java/com/databricks/sdk/core/oauth/Consent.java | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 73df7cd29..f2c45e9c4 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -388,6 +388,11 @@ private void sendSuccess(HttpExchange exchange) throws IOException { exchange.close(); } + /** + * Wait and return the params. + * + * This method might throw an exception in case of timeout. + */ public Map getParams() { synchronized (lock) { if (params == null) { From 3a54783be4f67538a6221d2923784f378d839b46 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 8 Aug 2025 08:38:39 +0000 Subject: [PATCH 9/9] Fix comment --- .../main/java/com/databricks/sdk/core/oauth/Consent.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index f2c45e9c4..19619d127 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -388,10 +388,10 @@ private void sendSuccess(HttpExchange exchange) throws IOException { exchange.close(); } - /** - * Wait and return the params. - * - * This method might throw an exception in case of timeout. + /** + * Wait and return the params. + * + *

This method might throw an exception in case of timeout. */ public Map getParams() { synchronized (lock) {