Better error message when private link enabled workspaces reject requests

apierr/errors.go

@@ -160,6 +160,11 @@ func GetAPIError(ctx context.Context, resp common.ResponseWrapper) error {
 		applyOverrides(ctx, apiError, resp.Response)
 		return apiError
 	}
+	// Attempts to access private link workspaces are redirected to the login page with a specific query parameter.
+	requestUrl := resp.Response.Request.URL
+	if isPrivateLinkRedirect(requestUrl) {
+		return privateLinkValidationError(requestUrl)
+	}
 	return nil
 }
 

apierr/errors_test.go

@@ -56,3 +56,20 @@ func TestApiErrorTransientRegexMatches(t *testing.T) {
 	ctx := context.Background()
 	assert.True(t, err.IsRetriable(ctx))
 }
+
+func TestApiErrorMapsPrivateLinkRedirect(t *testing.T) {
+	resp := common.ResponseWrapper{
+		Response: &http.Response{
+			Request: &http.Request{
+				URL: &url.URL{
+					Host:     "adb-12345678910.1.azuredatabricks.net",
+					Path:     "/login.html",
+					RawQuery: "error=private-link-validation-error",
+				},
+			},
+		},
+	}
+	err := GetAPIError(context.Background(), resp)
+	assert.ErrorIs(t, err, ErrPermissionDenied)
+	assert.Equal(t, err.(*APIError).ErrorCode, "PRIVATE_LINK_VALIDATION_ERROR")
+}

apierr/private_link.go

@@ -0,0 +1,73 @@
+package apierr
+
+import (
+	"fmt"
+	"net/url"
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go/common/environment"
+)
+
+// Metadata about the private link product. Private link redirects users to the
+// login page with a query parameter that indicates the error. This struct
+// contains information about the private link service, the endpoint name, and a
+// reference page for more information.
+//
+// Eventually, the REST API should return an error directly when a request is
+// made from a network that does not have access to the workspace. Once that
+// happens, this struct can be removed.
+type privateLinkInfo struct {
+	// The name of the private link service (e.g. AWS PrivateLink, Azure Private
+	// Link, etc.)
+	serviceName string
+
+	// The name of the private link endpoint (e.g. AWS VPC endpoint, Azure Private
+	// Link endpoint, etc.)
+	endpointName string
+
+	// A reference page for more information about the private link service.
+	referencePage string
+}
+
+func (p privateLinkInfo) errorMessage() string {
+	privateLinkValidationError := fmt.Sprintf(
+		`The requested workspace has %[1]s enabled and is not accessible from
+the current network. Ensure that %[1]s is properly configured and that your
+device has access to the %s. For more information, see %s.`,
+		p.serviceName, p.endpointName, p.referencePage)
+	return strings.ReplaceAll(privateLinkValidationError, "\n", " ")
+}
+
+// Map of private link information by cloud.
+var privateLinkInfoMap = map[environment.Cloud]privateLinkInfo{
+	environment.CloudAWS: {
+		serviceName:   "AWS PrivateLink",
+		endpointName:  "AWS VPC endpoint",
+		referencePage: "https://docs.databricks.com/en/security/network/classic/privatelink.html",
+	},
+	environment.CloudAzure: {
+		serviceName:   "Azure Private Link",
+		endpointName:  "Azure Private Link endpoint",
+		referencePage: "https://learn.microsoft.com/en-us/azure/databricks/security/network/classic/private-link-standard#authentication-troubleshooting",
+	},
+	environment.CloudGCP: {
+		serviceName:   "Private Service Connect",
+		endpointName:  "GCP VPC endpoint",
+		referencePage: "https://docs.gcp.databricks.com/en/security/network/classic/private-service-connect.html",
+	},
+}
+
+func isPrivateLinkRedirect(url *url.URL) bool {
+	return strings.Contains(url.RawQuery, "error=private-link-validation-error") && url.EscapedPath() == "/login.html"
+}
+
+func privateLinkValidationError(url *url.URL) *APIError {
+	env := environment.GetEnvironmentForHostname(url.Host)
+	info := privateLinkInfoMap[env.Cloud]
+	return &APIError{
+		ErrorCode:  "PRIVATE_LINK_VALIDATION_ERROR",
+		StatusCode: 403,
+		Message:    info.errorMessage(),
+		unwrap:     ErrPermissionDenied,
+	}
+}

common/environment/azure.go

@@ -0,0 +1,32 @@
+package environment
+
+type azureEnvironment struct {
+	Name                      string `json:"name"`
+	ServiceManagementEndpoint string `json:"serviceManagementEndpoint"`
+	ResourceManagerEndpoint   string `json:"resourceManagerEndpoint"`
+	ActiveDirectoryEndpoint   string `json:"activeDirectoryEndpoint"`
+}
+
+// based on github.com/Azure/go-autorest/autorest/azure/azureEnvironments.go
+var (
+	AzurePublicCloud = azureEnvironment{
+		Name:                      "PUBLIC",
+		ServiceManagementEndpoint: "https://management.core.windows.net/",
+		ResourceManagerEndpoint:   "https://management.azure.com/",
+		ActiveDirectoryEndpoint:   "https://login.microsoftonline.com/",
+	}
+
+	AzureUsGovernmentCloud = azureEnvironment{
+		Name:                      "USGOVERNMENT",
+		ServiceManagementEndpoint: "https://management.core.usgovcloudapi.net/",
+		ResourceManagerEndpoint:   "https://management.usgovcloudapi.net/",
+		ActiveDirectoryEndpoint:   "https://login.microsoftonline.us/",
+	}
+
+	AzureChinaCloud = azureEnvironment{
+		Name:                      "CHINA",
+		ServiceManagementEndpoint: "https://management.core.chinacloudapi.cn/",
+		ResourceManagerEndpoint:   "https://management.chinacloudapi.cn/",
+		ActiveDirectoryEndpoint:   "https://login.chinacloudapi.cn/",
+	}
+)

common/environment/environments.go

@@ -0,0 +1,85 @@
+package environment
+
+import (
+	"fmt"
+	"strings"
+)
+
+type Cloud string
+
+const (
+	CloudAWS   Cloud = "AWS"
+	CloudAzure Cloud = "Azure"
+	CloudGCP   Cloud = "GCP"
+)
+
+type DatabricksEnvironment struct {
+	Cloud              Cloud
+	DnsZone            string
+	AzureApplicationID string
+	AzureEnvironment   *azureEnvironment
+}
+
+func (de DatabricksEnvironment) DeploymentURL(name string) string {
+	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
+func DefaultEnvironment() DatabricksEnvironment {
+	return DatabricksEnvironment{
+		Cloud:   CloudAWS,
+		DnsZone: ".cloud.databricks.com",
+	}
+
+}
+
+var envs = []DatabricksEnvironment{
+	{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: &AzurePublicCloud},
+	{Cloud: CloudAzure, DnsZone: ".staging.azuredatabricks.net", AzureApplicationID: "4a67d088-db5c-48f1-9ff2-0aace800ae68", AzureEnvironment: &AzurePublicCloud},
+	{Cloud: CloudAzure, DnsZone: ".azuredatabricks.net", AzureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", AzureEnvironment: &AzurePublicCloud},
+	{Cloud: CloudAzure, DnsZone: ".databricks.azure.us", AzureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", AzureEnvironment: &AzureUsGovernmentCloud},
+	{Cloud: CloudAzure, DnsZone: ".databricks.azure.cn", AzureApplicationID: "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d", AzureEnvironment: &AzureChinaCloud},
+
+	{Cloud: CloudGCP, DnsZone: ".dev.gcp.databricks.com"},
+	{Cloud: CloudGCP, DnsZone: ".staging.gcp.databricks.com"},
+	{Cloud: CloudGCP, DnsZone: ".gcp.databricks.com"},
+}
+
+func AllEnvironments() []DatabricksEnvironment {
+	return append([]DatabricksEnvironment{}, envs...)
+}
+
+func GetEnvironmentForHostname(hostname string) DatabricksEnvironment {
+	for _, e := range envs {
+		if strings.HasSuffix(hostname, e.DnsZone) {
+			return e
+		}
+	}
+	return DefaultEnvironment()
+}

common/environment/environments_test.go

@@ -0,0 +1,11 @@
+package environment
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDefaultEnvironmentIsReturned(t *testing.T) {
+	assert.Equal(t, ".cloud.databricks.com", GetEnvironmentForHostname("").DnsZone)
+}

config/auth_permutations_test.go

@@ -113,14 +113,6 @@ func (cf configFixture) configureProviderAndReturnConfig(t *testing.T) (*Config,
 		AzureResourceID:   cf.AzureResourceID,
 		AuthType:          cf.AuthType,
 	}
-	if client.IsAzure() {
-		client.DatabricksEnvironment = &DatabricksEnvironment{
-			Cloud:              CloudAzure,
-			DnsZone:            cf.Host,
-			AzureApplicationID: "abc",
-			AzureEnvironment:   &AzurePublicCloud,
-		}
-	}
 	err := client.Authenticate(&http.Request{Header: http.Header{}})
 	if err != nil {
 		return nil, err
@@ -369,14 +361,14 @@ func TestConfig_AzurePAT(t *testing.T) {
 
 func TestConfig_AzureCliHost(t *testing.T) {
 	configFixture{
-		Host:            "x",          // adb-123.4.azuredatabricks.net
+		Host:            "https://adb-123.4.azuredatabricks.net",
 		AzureResourceID: azResourceID, // skips ensureWorkspaceUrl
 		Env: map[string]string{
 			"PATH": testdataPath(),
 			"HOME": "testdata/azure",
 		},
 		AssertAzure: true,
-		AssertHost:  "https://x",
+		AssertHost:  "https://adb-123.4.azuredatabricks.net",
 		AssertAuth:  "azure-cli",
 	}.apply(t)
 }
@@ -421,13 +413,13 @@ func TestConfig_AzureCliHostAndResourceID(t *testing.T) {
 	configFixture{
 		// omit request to management endpoint to get workspace properties
 		AzureResourceID: azResourceID,
-		Host:            "x",
+		Host:            "https://adb-123.4.azuredatabricks.net",
 		Env: map[string]string{
 			"PATH": testdataPath(),
 			"HOME": "testdata", // .databrickscfg has DEFAULT profile
 		},
 		AssertAzure: true,
-		AssertHost:  "https://x",
+		AssertHost:  "https://adb-123.4.azuredatabricks.net",
 		AssertAuth:  "azure-cli",
 	}.apply(t)
 }
@@ -436,28 +428,28 @@ func TestConfig_AzureCliHostAndResourceID_ConfigurationPrecedence(t *testing.T)
 	configFixture{
 		// omit request to management endpoint to get workspace properties
 		AzureResourceID: azResourceID,
-		Host:            "x",
+		Host:            "https://adb-123.4.azuredatabricks.net",
 		Env: map[string]string{
 			"PATH":                      testdataPath(),
 			"HOME":                      "testdata/azure",
 			"DATABRICKS_CONFIG_PROFILE": "justhost",
 		},
 		AssertAzure: true,
-		AssertHost:  "https://x",
+		AssertHost:  "https://adb-123.4.azuredatabricks.net",
 		AssertAuth:  "azure-cli",
 	}.apply(t)
 }
 
 func TestConfig_AzureAndPasswordConflict(t *testing.T) { // TODO: this breaks
 	configFixture{
-		Host:            "x",
+		Host:            "https://adb-123.4.azuredatabricks.net",
 		AzureResourceID: azResourceID,
 		Env: map[string]string{
 			"PATH":                testdataPath(),
 			"HOME":                "testdata/azure",
 			"DATABRICKS_USERNAME": "x",
 		},
-		AssertError: "validate: more than one authorization method configured: azure and basic. Config: host=x, username=x, azure_workspace_resource_id=/sub/rg/ws. Env: DATABRICKS_USERNAME",
+		AssertError: "validate: more than one authorization method configured: azure and basic. Config: host=https://adb-123.4.azuredatabricks.net, username=x, azure_workspace_resource_id=/sub/rg/ws. Env: DATABRICKS_USERNAME",
 	}.apply(t)
 }
 

config/azure.go

@@ -110,37 +110,6 @@ func (c *Config) mapAzureResourceManagerError(defaultErr *httpclient.HttpError)
 	}
 }
 
-type azureEnvironment struct {
-	Name                      string `json:"name"`
-	ServiceManagementEndpoint string `json:"serviceManagementEndpoint"`
-	ResourceManagerEndpoint   string `json:"resourceManagerEndpoint"`
-	ActiveDirectoryEndpoint   string `json:"activeDirectoryEndpoint"`
-}
-
-// based on github.com/Azure/go-autorest/autorest/azure/azureEnvironments.go
-var (
-	AzurePublicCloud = azureEnvironment{
-		Name:                      "PUBLIC",
-		ServiceManagementEndpoint: "https://management.core.windows.net/",
-		ResourceManagerEndpoint:   "https://management.azure.com/",
-		ActiveDirectoryEndpoint:   "https://login.microsoftonline.com/",
-	}
-
-	AzureUsGovernmentCloud = azureEnvironment{
-		Name:                      "USGOVERNMENT",
-		ServiceManagementEndpoint: "https://management.core.usgovcloudapi.net/",
-		ResourceManagerEndpoint:   "https://management.usgovcloudapi.net/",
-		ActiveDirectoryEndpoint:   "https://login.microsoftonline.us/",
-	}
-
-	AzureChinaCloud = azureEnvironment{
-		Name:                      "CHINA",
-		ServiceManagementEndpoint: "https://management.core.chinacloudapi.cn/",
-		ResourceManagerEndpoint:   "https://management.chinacloudapi.cn/",
-		ActiveDirectoryEndpoint:   "https://login.chinacloudapi.cn/",
-	}
-)
-
 type azureHostResolver interface {
 	tokenSourceFor(ctx context.Context, cfg *Config, aadEndpoint, resource string) oauth2.TokenSource
 }

config/config.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sdk-go/common"
+	"github.com/databricks/databricks-sdk-go/common/environment"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/logger"
 )
@@ -119,7 +120,7 @@ type Config struct {
 	HTTPTransport http.RoundTripper
 
 	// Environment override to return when resolving the current environment.
-	DatabricksEnvironment *DatabricksEnvironment
+	DatabricksEnvironment *environment.DatabricksEnvironment
 
 	Loaders []Loader
 
@@ -212,12 +213,12 @@ func (c *Config) IsAzure() bool {
 	if c.AzureResourceID != "" {
 		return true
 	}
-	return c.Environment().Cloud == CloudAzure
+	return c.Environment().Cloud == environment.CloudAzure
 }
 
 // IsGcp returns if the client is configured for Databricks on Google Cloud.
 func (c *Config) IsGcp() bool {
-	return c.Environment().Cloud == CloudGCP
+	return c.Environment().Cloud == environment.CloudGCP
 }
 
 // IsAws returns if the client is configured for Databricks on AWS.