..

config/auth_azure_cli.go

@@ -44,7 +44,7 @@ func (c AzureCliCredentials) tokenSourceFor(
 // If the user can't access the service management endpoint, we assume they are in case 2 and do not pass the service
 // management token. Otherwise, we always pass the service management token.
 func (c AzureCliCredentials) getVisitor(ctx context.Context, cfg *Config, inner oauth2.TokenSource) (func(*http.Request) error, error) {
-	ts := &azureCliTokenSource{cfg.Environment().AzureEnvironment.ServiceManagementEndpoint, ""}
+	ts := &azureCliTokenSource{cfg.Environment().AzureServiceManagementEndpoint(), ""}
 	t, err := ts.Token()
 	if err != nil {
 		logger.Debugf(ctx, "Not including service management token in headers: %v", err)
@@ -59,7 +59,7 @@ func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (func(*
 		return nil, nil
 	}
 	// Eagerly get a token to fail fast in case the user is not logged in with the Azure CLI.
-	ts := &azureCliTokenSource{cfg.Environment().AzureApplicationID, cfg.AzureResourceID}
+	ts := &azureCliTokenSource{cfg.Environment().azureApplicationID, cfg.AzureResourceID}
 	t, err := ts.Token()
 	if err != nil {
 		if strings.Contains(err.Error(), "No subscription found") {

config/auth_azure_client_secret.go

@@ -51,9 +51,9 @@ func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config
 	logger.Infof(ctx, "Generating AAD token for Service Principal (%s)", cfg.AzureClientID)
 	refreshCtx := context.Background()
 	env := cfg.Environment()
-	aadEndpoint := env.AzureEnvironment.ActiveDirectoryEndpoint
-	managementEndpoint := env.AzureEnvironment.ServiceManagementEndpoint
-	inner := azureReuseTokenSource(nil, c.tokenSourceFor(refreshCtx, cfg, aadEndpoint, env.AzureApplicationID))
+	aadEndpoint := env.AzureActiveDirectoryEndpoint()
+	managementEndpoint := env.AzureServiceManagementEndpoint()
+	inner := azureReuseTokenSource(nil, c.tokenSourceFor(refreshCtx, cfg, aadEndpoint, env.azureApplicationID))
 	management := azureReuseTokenSource(nil, c.tokenSourceFor(refreshCtx, cfg, aadEndpoint, managementEndpoint))
 	return azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken)), nil
 }

config/auth_azure_msi.go

@@ -41,11 +41,11 @@ func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (func(*
 	}
 	logger.Debugf(ctx, "Generating AAD token via Azure MSI")
 	inner := azureReuseTokenSource(nil, azureMsiTokenSource{
-		resource: env.AzureApplicationID,
+		resource: env.azureApplicationID,
 		clientId: cfg.AzureClientID,
 	})
 	management := azureReuseTokenSource(nil, azureMsiTokenSource{
-		resource: env.AzureEnvironment.ServiceManagementEndpoint,
+		resource: env.AzureServiceManagementEndpoint(),
 		clientId: cfg.AzureClientID,
 	})
 	return azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken)), nil

config/auth_permutations_test.go

@@ -116,9 +116,9 @@ func (cf configFixture) configureProviderAndReturnConfig(t *testing.T) (*Config,
 	if client.IsAzure() {
 		client.DatabricksEnvironments = append(client.DatabricksEnvironments, DatabricksEnvironment{
 			Cloud:              CloudAzure,
-			DnsZone:            cf.Host,
-			AzureApplicationID: "abc",
-			AzureEnvironment:   &publicCloud,
+			dnsZone:            cf.Host,
+			azureApplicationID: "abc",
+			azureEnvironment:   &publicCloud,
 		})
 	}
 	err := client.Authenticate(&http.Request{Header: http.Header{}})

config/azure.go

@@ -48,7 +48,7 @@ func (c *Config) azureEnsureWorkspaceUrl(ctx context.Context, ahr azureHostResol
 	if c.AzureResourceID == "" || c.Host != "" {
 		return nil
 	}
-	azureEnv := c.Environment().AzureEnvironment
+	azureEnv := c.Environment().azureEnvironment
 	// azure resource ID can also be used in lieu of host by some of the clients, like Terraform
 	management := ahr.tokenSourceFor(ctx, c, azureEnv.ActiveDirectoryEndpoint, azureEnv.ResourceManagerEndpoint)
 	var workspaceMetadata struct {

config/environments.go

@@ -16,37 +16,59 @@ const (
 
 type DatabricksEnvironment struct {
 	Cloud              Cloud
-	DnsZone            string
-	AzureApplicationID string
-	AzureEnvironment   *azureEnvironment
+	dnsZone            string
+	azureApplicationID string
+	azureEnvironment   *azureEnvironment
 }
 
 func (de DatabricksEnvironment) DeploymentURL(name string) string {
-	return fmt.Sprintf("https://%s%s", name, de.DnsZone)
+	return fmt.Sprintf("https://%s%s", name, de.dnsZone)
 }
 
+func (de DatabricksEnvironment) AzureServiceManagementEndpoint() string {
+	if de.azureEnvironment == nil {
+		return ""
+	}
+	return de.azureEnvironment.ServiceManagementEndpoint
+}
+
+func (de DatabricksEnvironment) AzureResourceManagerEndpoint() string {
+	if de.azureEnvironment == nil {
+		return ""
+	}
+	return de.azureEnvironment.ResourceManagerEndpoint
+}
+
+func (de DatabricksEnvironment) AzureActiveDirectoryEndpoint() string {
+	if de.azureEnvironment == nil {
+		return ""
+	}
+	return de.azureEnvironment.ActiveDirectoryEndpoint
+}
+
+// we default to AWS Prod environment since this case will be a hit for PVC
 var defaultEnvironment = DatabricksEnvironment{
 	Cloud:   CloudAWS,
-	DnsZone: ".cloud.databricks.com",
+	dnsZone: ".cloud.databricks.com",
 }
 
 var envs = []DatabricksEnvironment{
-	{Cloud: CloudUnspecified, DnsZone: "localhost"},
+	{Cloud: CloudUnspecified, dnsZone: "localhost"},
 
-	{Cloud: CloudAWS, DnsZone: ".dev.databricks.com"},
-	{Cloud: CloudAWS, DnsZone: ".staging.cloud.databricks.com"},
-	{Cloud: CloudAWS, DnsZone: ".cloud.databricks.us"},
+	{Cloud: CloudAWS, dnsZone: ".dev.databricks.com"},
+	{Cloud: CloudAWS, dnsZone: ".staging.cloud.databricks.com"},
+	{Cloud: CloudAWS, dnsZone: ".cloud.databricks.us"},
 	defaultEnvironment,
 
-	{Cloud: CloudAzure, DnsZone: ".dev.azuredatabricks.net", AzureApplicationID: "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", AzureEnvironment: &publicCloud},
-	{Cloud: CloudAzure, DnsZone: ".staging.azuredatabricks.net", AzureApplicationID: "4a67d088-db5c-48f1-9ff2-0aace800ae68", AzureEnvironment: &publicCloud},
-	{Cloud: CloudAzure, DnsZone: ".azuredatabricks.net", AzureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", AzureEnvironment: &publicCloud},
-	{Cloud: CloudAzure, DnsZone: ".databricks.azure.us", AzureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", AzureEnvironment: &usGovernmentCloud},
-	{Cloud: CloudAzure, DnsZone: ".databricks.azure.cn", AzureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", AzureEnvironment: &chinaCloud},
+	{Cloud: CloudAzure, dnsZone: ".dev.azuredatabricks.net", azureApplicationID: "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc", azureEnvironment: &publicCloud},
+	{Cloud: CloudAzure, dnsZone: ".staging.azuredatabricks.net", azureApplicationID: "4a67d088-db5c-48f1-9ff2-0aace800ae68", azureEnvironment: &publicCloud},
+	{Cloud: CloudAzure, dnsZone: ".azuredatabricks.net", azureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", azureEnvironment: &publicCloud},
+	{Cloud: CloudAzure, dnsZone: ".databricks.azure.us", azureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", azureEnvironment: &usGovernmentCloud},
+	{Cloud: CloudAzure, dnsZone: ".databricks.azure.cn", azureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", azureEnvironment: &chinaCloud},
 
-	{Cloud: CloudGCP, DnsZone: ".dev.gcp.databricks.com"},
-	{Cloud: CloudGCP, DnsZone: ".staging.gcp.databricks.com"},
-	{Cloud: CloudGCP, DnsZone: ".gcp.databricks.com"},
+	{Cloud: CloudGCP, dnsZone: ".dev.gcp.databricks.com"},
+	{Cloud: CloudGCP, dnsZone: ".staging.gcp.databricks.com"},
+	{Cloud: CloudGCP, dnsZone: ".gcp.databricks.com"},
 }
 
 func (c *Config) Environment() DatabricksEnvironment {
@@ -61,15 +83,15 @@ func (c *Config) Environment() DatabricksEnvironment {
 			if v.Cloud != CloudAzure {
 				continue
 			}
-			if v.AzureEnvironment.Name != azureEnv {
+			if v.azureEnvironment.Name != azureEnv {
 				continue
 			}
 			return v
 		}
 	}
 	hostname := c.CanonicalHostName()
 	for _, e := range append(c.DatabricksEnvironments, envs...) {
-		if strings.HasSuffix(hostname, e.DnsZone) {
+		if strings.HasSuffix(hostname, e.dnsZone) {
 			return e
 		}
 	}