From 94b78cac6695708e50c67b6ddabe52db5bc5b6ba Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Fri, 17 May 2024 11:01:21 +0200
Subject: [PATCH] work

---
 apierr/errors.go            |  5 +++--
 apierr/errors_test.go       |  5 +++--
 apierr/private_link.go      | 26 +++++++++++++++++++++++---
 common/environment/azure.go |  3 +--
 4 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/apierr/errors.go b/apierr/errors.go
index 2d6b602d0..1f38e0962 100644
--- a/apierr/errors.go
+++ b/apierr/errors.go
@@ -161,8 +161,9 @@ func GetAPIError(ctx context.Context, resp common.ResponseWrapper) error {
 		return apiError
 	}
 	// Attempts to access private link workspaces are redirected to the login page with a specific query parameter.
-	if strings.Contains(resp.Response.Request.URL.RawQuery, "error=private-link-validation-error") {
-		return PrivateLinkValidationError(resp.Response.Request.URL)
+	requestUrl := resp.Response.Request.URL
+	if isPrivateLinkRedirect(requestUrl) {
+		return privateLinkValidationError(requestUrl)
 	}
 	return nil
 }
diff --git a/apierr/errors_test.go b/apierr/errors_test.go
index 089954ba5..be725164f 100644
--- a/apierr/errors_test.go
+++ b/apierr/errors_test.go
@@ -62,7 +62,8 @@ func TestApiErrorMapsPrivateLinkRedirect(t *testing.T) {
 		Response: &http.Response{
 			Request: &http.Request{
 				URL: &url.URL{
-					Host: "adb-12345678910.1.azuredatabricks.net",
+					Host:     "adb-12345678910.1.azuredatabricks.net",
+					Path:     "/login.html",
 					RawQuery: "error=private-link-validation-error",
 				},
 			},
@@ -71,4 +72,4 @@ func TestApiErrorMapsPrivateLinkRedirect(t *testing.T) {
 	err := GetAPIError(context.Background(), resp)
 	assert.ErrorIs(t, err, ErrPermissionDenied)
 	assert.Equal(t, err.(*APIError).ErrorCode, "PRIVATE_LINK_VALIDATION_ERROR")
-}
\ No newline at end of file
+}
diff --git a/apierr/private_link.go b/apierr/private_link.go
index e1a398f59..16ce19ebc 100644
--- a/apierr/private_link.go
+++ b/apierr/private_link.go
@@ -8,9 +8,24 @@ import (
 	"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 {
-	serviceName   string
-	endpointName  string
+	// 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
 }
 
@@ -23,6 +38,7 @@ device has access to the %s. For more information, see %s.`,
 	return strings.ReplaceAll(privateLinkValidationError, "\n", " ")
 }
 
+// Map of private link information by cloud.
 var privateLinkInfoMap = map[environment.Cloud]privateLinkInfo{
 	environment.CloudAWS: {
 		serviceName:   "AWS PrivateLink",
@@ -41,7 +57,11 @@ var privateLinkInfoMap = map[environment.Cloud]privateLinkInfo{
 	},
 }
 
-func PrivateLinkValidationError(url *url.URL) *APIError {
+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{
diff --git a/common/environment/azure.go b/common/environment/azure.go
index 5793e66e8..cfda1028b 100644
--- a/common/environment/azure.go
+++ b/common/environment/azure.go
@@ -1,6 +1,5 @@
 package environment
 
-
 type azureEnvironment struct {
 	Name                      string `json:"name"`
 	ServiceManagementEndpoint string `json:"serviceManagementEndpoint"`
@@ -30,4 +29,4 @@ var (
 		ResourceManagerEndpoint:   "https://management.chinacloudapi.cn/",
 		ActiveDirectoryEndpoint:   "https://login.chinacloudapi.cn/",
 	}
-)
\ No newline at end of file
+)