From c118ad433ff6bb94dfd8ab8984c11c5642180f6f Mon Sep 17 00:00:00 2001 From: Jacky Hu Date: Tue, 23 Jan 2024 11:02:19 -0800 Subject: [PATCH] [PECO-1414] Change AAD openid config url to Databirck workspace Signed-off-by: Jacky Hu --- src/databricks/sql/auth/endpoint.py | 5 ++++- tests/unit/test_endpoint.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/databricks/sql/auth/endpoint.py b/src/databricks/sql/auth/endpoint.py index c0ce0f9db..34ebe4234 100644 --- a/src/databricks/sql/auth/endpoint.py +++ b/src/databricks/sql/auth/endpoint.py @@ -91,7 +91,10 @@ def get_authorization_url(self, hostname: str): return f"{get_databricks_oidc_url(hostname)}/oauth2/v2.0/authorize" def get_openid_config_url(self, hostname: str): - return "https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration" + # This well known openid url wil return Azure well-known openid configuration response + # or Databricks OAuth well-known URL + idp_url = get_databricks_oidc_url(hostname) + return f"{idp_url}/.well-known/openid-configuration" class AwsOAuthEndpointCollection(OAuthEndpointCollection): diff --git a/tests/unit/test_endpoint.py b/tests/unit/test_endpoint.py index 63393039b..8e5bc606c 100644 --- a/tests/unit/test_endpoint.py +++ b/tests/unit/test_endpoint.py @@ -36,7 +36,7 @@ def test_oauth_endpoint(self): CloudType.AZURE, azure_host, f"https://{azure_host}/oidc/oauth2/v2.0/authorize", - "https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration", + f"https://{azure_host}/oidc/.well-known/openid-configuration", [azure_scope, "offline_access"], [azure_scope] )]