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 23fd990b9..60621607d 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 @@ -798,15 +798,18 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException { } if (isAzure() && getAzureClientId() != null) { - Request request = new Request("GET", getHost() + "/oidc/oauth2/v2.0/authorize"); - request.setRedirectionBehavior(false); - Response resp = getHttpClient().execute(request); - String realAuthUrl = resp.getFirstHeader("location"); - if (realAuthUrl == null) { - return null; + ApiClient apiClient = + new ApiClient.Builder() + .withHttpClient(getHttpClient()) + .withGetHostFunc(v -> getHost()) + .build(); + try { + return apiClient.execute( + new Request("GET", "/oidc/.well-known/oauth-authorization-server"), + OpenIDConnectEndpoints.class); + } catch (IOException e) { + throw new DatabricksException("IO error: " + e.getMessage(), e); } - return new OpenIDConnectEndpoints( - realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl); } if (isAccountClient() && getAccountId() != null) { String prefix = getHost() + "/oidc/accounts/" + getAccountId(); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcAccountEndpointIT.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcAccountEndpointIT.java new file mode 100644 index 000000000..810a92d25 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcAccountEndpointIT.java @@ -0,0 +1,62 @@ +package com.databricks.sdk.integration; + +import com.databricks.sdk.AccountClient; +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.integration.framework.EnvContext; +import com.databricks.sdk.integration.framework.EnvTest; +import java.io.IOException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.junit.jupiter.api.extension.ExtendWith; + +@EnvContext("account") +@ExtendWith(EnvTest.class) +public class OidcAccountEndpointIT { + @Test + void checkEndpoints(AccountClient a) throws IOException { + OpenIDConnectEndpoints endpoints = a.config().getOidcEndpoints(); + String host = a.config().getHost(); + String accountId = a.config().getAccountId(); + if (a.config().isAzure()) { + assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token"); + } else { + assert endpoints + .getAuthorizationEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize"); + assert endpoints + .getTokenEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/token"); + } + } + + @Test + void unifiedEndpoints(AccountClient a) throws IOException { + DatabricksConfig cfg = a.config(); + cfg.setExperimentalIsUnifiedHost(true); + OpenIDConnectEndpoints endpoints = cfg.getOidcEndpoints(); + String host = cfg.getHost(); + String accountId = cfg.getAccountId(); + assert endpoints + .getAuthorizationEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/accounts/" + accountId + "/v1/token"); + } + + @Test + @EnabledIfEnvironmentVariable(named = "UNIFIED_HOST", matches = ".+") + void unifiedEndpointsForSpog(AccountClient a) throws IOException { + String unifiedHost = System.getenv("UNIFIED_HOST"); + DatabricksConfig cfg = a.config(); + cfg.setHost(unifiedHost); + cfg.setExperimentalIsUnifiedHost(true); + OpenIDConnectEndpoints endpoints = cfg.getOidcEndpoints(); + String host = cfg.getHost(); + String accountId = cfg.getAccountId(); + assert endpoints + .getAuthorizationEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/accounts/" + accountId + "/v1/token"); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcAccountEndpointIT.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcAccountEndpointIT.java new file mode 100644 index 000000000..edb6f6da6 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcAccountEndpointIT.java @@ -0,0 +1,49 @@ +package com.databricks.sdk.integration; + +import com.databricks.sdk.AccountClient; +import com.databricks.sdk.core.DatabricksConfig; +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.integration.framework.EnvContext; +import com.databricks.sdk.integration.framework.EnvTest; +import java.io.IOException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.junit.jupiter.api.extension.ExtendWith; + +@EnvContext("ucacct") +@ExtendWith(EnvTest.class) +public class OidcUcAccountEndpointIT { + @Test + void checkEndpoints(AccountClient a) throws IOException { + OpenIDConnectEndpoints endpoints = a.config().getOidcEndpoints(); + String host = a.config().getHost(); + String accountId = a.config().getAccountId(); + if (a.config().isAzure()) { + assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token"); + } else { + assert endpoints + .getAuthorizationEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize"); + assert endpoints + .getTokenEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/token"); + } + } + + @Test + @EnabledIfEnvironmentVariable(named = "UNIFIED_HOST", matches = ".+") + void unifiedEndpointsForSpog(AccountClient a) throws IOException { + String unifiedHost = System.getenv("UNIFIED_HOST"); + DatabricksConfig cfg = a.config(); + cfg.setHost(unifiedHost); + cfg.setExperimentalIsUnifiedHost(true); + OpenIDConnectEndpoints endpoints = cfg.getOidcEndpoints(); + String host = cfg.getHost(); + String accountId = cfg.getAccountId(); + assert endpoints + .getAuthorizationEndpoint() + .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/accounts/" + accountId + "/v1/token"); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcWorkspaceEndpointIT.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcWorkspaceEndpointIT.java new file mode 100644 index 000000000..6f7b5709e --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcWorkspaceEndpointIT.java @@ -0,0 +1,21 @@ +package com.databricks.sdk.integration; + +import com.databricks.sdk.WorkspaceClient; +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.integration.framework.EnvContext; +import com.databricks.sdk.integration.framework.EnvTest; +import java.io.IOException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; + +@EnvContext("ucws") +@ExtendWith(EnvTest.class) +public class OidcUcWorkspaceEndpointIT { + @Test + void checkEndpoints(WorkspaceClient w) throws IOException { + OpenIDConnectEndpoints endpoints = w.config().getOidcEndpoints(); + String host = w.config().getHost(); + assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token"); + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcWorkspaceEndpointIT.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcWorkspaceEndpointIT.java new file mode 100644 index 000000000..800411473 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcWorkspaceEndpointIT.java @@ -0,0 +1,21 @@ +package com.databricks.sdk.integration; + +import com.databricks.sdk.WorkspaceClient; +import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints; +import com.databricks.sdk.integration.framework.EnvContext; +import com.databricks.sdk.integration.framework.EnvTest; +import java.io.IOException; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; + +@EnvContext("workspace") +@ExtendWith(EnvTest.class) +public class OidcWorkspaceEndpointIT { + @Test + void checkEndpoints(WorkspaceClient w) throws IOException { + OpenIDConnectEndpoints endpoints = w.config().getOidcEndpoints(); + String host = w.config().getHost(); + assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize"); + assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token"); + } +}