From 04b9dc6dc941fef97fd02b75b8f7fe2802cf1214 Mon Sep 17 00:00:00 2001 From: Erfan Mahmoodnejad Date: Thu, 11 Apr 2024 21:27:13 +1200 Subject: [PATCH 1/5] Fixed OAuth for Azure Databricks --- auth/oauth/oauth.go | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index ad76dad..0e7ab43 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -10,9 +10,18 @@ import ( "golang.org/x/oauth2" ) -const ( - azureTenantId = "4a67d088-db5c-48f1-9ff2-0aace800ae68" -) +type AzureTenant struct { + DnsZone string + AzureApplicationID string +} + +var azureTenants = map[string]string{ + ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", + ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68", + ".azuredatabricks.net": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", + ".databricks.azure.us": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", + ".databricks.azure.cn": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", +} func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error) { if ctx == nil { @@ -52,7 +61,7 @@ func GetScopes(hostName string, scopes []string) []string { cloudType := InferCloudFromHost(hostName) if cloudType == Azure { - userImpersonationScope := fmt.Sprintf("%s/user_impersonation", azureTenantId) + userImpersonationScope := fmt.Sprintf("%s/user_impersonation", azureTenants[GetAzureDnsZone(hostName)]) if !HasScope(scopes, userImpersonationScope) { scopes = append(scopes, userImpersonationScope) } @@ -133,3 +142,12 @@ func InferCloudFromHost(hostname string) CloudType { } return Unknown } + +func GetAzureDnsZone(hostname string) string { + for _, d := range databricksAzureDomains { + if strings.Contains(hostname, d) { + return d + } + } + return "" +} From f50a9b002f0777514feac67e80442780addc7b77 Mon Sep 17 00:00:00 2001 From: Erfan Mahmoodnejad Date: Thu, 11 Apr 2024 21:27:13 +1200 Subject: [PATCH 2/5] Fixed OAuth for Azure Databricks Signed-off-by: Erfan Mahmoodnejad --- auth/oauth/oauth.go | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index ad76dad..0e7ab43 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -10,9 +10,18 @@ import ( "golang.org/x/oauth2" ) -const ( - azureTenantId = "4a67d088-db5c-48f1-9ff2-0aace800ae68" -) +type AzureTenant struct { + DnsZone string + AzureApplicationID string +} + +var azureTenants = map[string]string{ + ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", + ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68", + ".azuredatabricks.net": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", + ".databricks.azure.us": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", + ".databricks.azure.cn": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", +} func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error) { if ctx == nil { @@ -52,7 +61,7 @@ func GetScopes(hostName string, scopes []string) []string { cloudType := InferCloudFromHost(hostName) if cloudType == Azure { - userImpersonationScope := fmt.Sprintf("%s/user_impersonation", azureTenantId) + userImpersonationScope := fmt.Sprintf("%s/user_impersonation", azureTenants[GetAzureDnsZone(hostName)]) if !HasScope(scopes, userImpersonationScope) { scopes = append(scopes, userImpersonationScope) } @@ -133,3 +142,12 @@ func InferCloudFromHost(hostname string) CloudType { } return Unknown } + +func GetAzureDnsZone(hostname string) string { + for _, d := range databricksAzureDomains { + if strings.Contains(hostname, d) { + return d + } + } + return "" +} From 51d1b5022a44086b7704b862ad3e7126921572c4 Mon Sep 17 00:00:00 2001 From: Erfan Mahmoodnejad Date: Thu, 11 Apr 2024 21:52:58 +1200 Subject: [PATCH 3/5] Removed the AzureTenant struct Signed-off-by: Erfan Mahmoodnejad --- auth/oauth/oauth.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index 0e7ab43..0df9d5c 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -10,11 +10,6 @@ import ( "golang.org/x/oauth2" ) -type AzureTenant struct { - DnsZone string - AzureApplicationID string -} - var azureTenants = map[string]string{ ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68", From 9af4f6543aff8012a3fe2b3835961e0179784d01 Mon Sep 17 00:00:00 2001 From: Erfan Mahmoodnejad Date: Thu, 11 Apr 2024 21:27:13 +1200 Subject: [PATCH 4/5] Fixed OAuth for Azure Databricks Signed-off-by: Erfan Mahmoodnejad --- auth/oauth/oauth.go | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index ad76dad..0e7ab43 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -10,9 +10,18 @@ import ( "golang.org/x/oauth2" ) -const ( - azureTenantId = "4a67d088-db5c-48f1-9ff2-0aace800ae68" -) +type AzureTenant struct { + DnsZone string + AzureApplicationID string +} + +var azureTenants = map[string]string{ + ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", + ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68", + ".azuredatabricks.net": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", + ".databricks.azure.us": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", + ".databricks.azure.cn": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", +} func GetEndpoint(ctx context.Context, hostName string) (oauth2.Endpoint, error) { if ctx == nil { @@ -52,7 +61,7 @@ func GetScopes(hostName string, scopes []string) []string { cloudType := InferCloudFromHost(hostName) if cloudType == Azure { - userImpersonationScope := fmt.Sprintf("%s/user_impersonation", azureTenantId) + userImpersonationScope := fmt.Sprintf("%s/user_impersonation", azureTenants[GetAzureDnsZone(hostName)]) if !HasScope(scopes, userImpersonationScope) { scopes = append(scopes, userImpersonationScope) } @@ -133,3 +142,12 @@ func InferCloudFromHost(hostname string) CloudType { } return Unknown } + +func GetAzureDnsZone(hostname string) string { + for _, d := range databricksAzureDomains { + if strings.Contains(hostname, d) { + return d + } + } + return "" +} From 1319cfbac0e3bddb3be3e066955c20b72ce3abea Mon Sep 17 00:00:00 2001 From: Erfan Mahmoodnejad Date: Thu, 11 Apr 2024 21:52:58 +1200 Subject: [PATCH 5/5] Removed the AzureTenant struct Signed-off-by: Erfan Mahmoodnejad --- auth/oauth/oauth.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index 0e7ab43..0df9d5c 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -10,11 +10,6 @@ import ( "golang.org/x/oauth2" ) -type AzureTenant struct { - DnsZone string - AzureApplicationID string -} - var azureTenants = map[string]string{ ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68",