databricks · mgyucht · Feb 1, 2024 · Jan 31, 2024 · Jan 31, 2024 · Jan 31, 2024
diff --git a/account_functions.go b/account_functions.go
@@ -0,0 +1,36 @@
+package databricks
+
+import "github.com/databricks/databricks-sdk-go/service/provisioning"
+
+// GetWorkspaceClient returns a WorkspaceClient for the given workspace. The
+// workspace can be fetched by calling w.Workspaces.Get() or w.Workspaces.List().
+//
+// The config used for the workspace is identical to that used for the account,
+// except that the host is set to the workspace host, and the account ID is
+// not set.
+//
+// Example:
+//
+//	a, err := databricks.NewAccountClient()
+//	if err != nil {
+//		panic(err)
+//	}
+//	ctx := context.Background()
+//	workspaces, err := a.Workspaces.List(ctx)
+//	if err != nil {
+//		panic(err)
+//	}
+//	w, err := a.GetWorkspaceClient(workspaces[0])
+//	if err != nil {
+//		panic(err)
+//	}
+//	me, err := w.CurrentUser.Me(ctx)
+func (c *AccountClient) GetWorkspaceClient(w provisioning.Workspace) (*WorkspaceClient, error) {
+	host := c.Config.Environment().DeploymentURL(w.DeploymentName)
+	cfg, err := c.Config.NewWithWorkspaceHost(host)
+	if err != nil {
+		return nil, err
+	}
+	cfg.AzureResourceID = w.AzureResourceId()
+	return NewWorkspaceClient((*Config)(cfg))
+}
diff --git a/config/config.go b/config/config.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -148,6 +149,48 @@ type Config struct {
 	auth func(r *http.Request) error
 }
 
+// NewWithWorkspaceHost returns a new instance of the Config with the host set to
+// the workspace host. Fields that are not relevant to workspace-level config,
+// like account ID, are omitted. Workspace-level attributes that cannot be
+// computed from the host alone, like Azure Resource ID, are also omitted.
+func (c *Config) NewWithWorkspaceHost(host string) (*Config, error) {
+	err := c.EnsureResolved()
+	if err != nil {
+		return nil, err
+	}
+
+	var fieldsToSkip = map[string]struct{}{
+		"Host":            {},
+		"AzureResourceID": {},
+		"AccountID":       {},
+	}
+	res := &Config{}
+	cv := reflect.ValueOf(c).Elem()
+	resv := reflect.ValueOf(res).Elem()
+	for i := 0; i < resv.NumField(); i++ {
+		field := resv.Field(i)
+		if !field.CanSet() {
+			continue
+		}
+		if _, ok := fieldsToSkip[resv.Type().Field(i).Name]; ok {
+			continue
+		}
+		field.Set(cv.Field(i))
+	}
+
+	res.Host = host
+	// We can reuse the same OAuth token refresh client and context. The
+	// reuseTokenSource internally locks.
+	res.refreshClient = c.refreshClient
+	res.refreshCtx = c.refreshCtx
+	// The config does not need to be re-resolved, as we reuse all attributes
+	// from the original config.
+	res.resolved = c.resolved
+	res.auth = c.auth
+	res.isTesting = c.isTesting
+	return res, nil
+}
+
 // Authenticate adds special headers to HTTP request to authorize it to work with Databricks REST API
 func (c *Config) Authenticate(r *http.Request) error {
 	err := c.EnsureResolved()

diff --git a/config/config_test.go b/config/config_test.go
@@ -29,3 +29,25 @@ func TestIsAccountClient_AwsWorkspace(t *testing.T) {
 	}
 	assert.False(t, c.IsAccountClient())
 }
+
+func TestNewWithWorkspaceHost(t *testing.T) {
+	c := &Config{
+		Host:               "https://accounts.cloud.databricks.com",
+		AccountID:          "123e4567-e89b-12d3-a456-426614174000",
+		AzureResourceID:    "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/test",
+		ClientID:           "client-id",
+		MetadataServiceURL: "http://",
+		resolved:           true,
+	}
+	c2, err := c.NewWithWorkspaceHost("https://my-workspace.cloud.databricks.us")
+	assert.NoError(t, err)
+	// Host should be updated
+	assert.Equal(t, "https://my-workspace.cloud.databricks.us", c2.Host)
+	// Account ID and Azure Resource ID should be cleared
+	assert.Equal(t, "", c2.AccountID)
+	assert.Equal(t, "", c2.AzureResourceID)
+	// Other fields should be preserved
+	assert.Equal(t, "client-id", c2.ClientID)
+	assert.Equal(t, "http://", c2.MetadataServiceURL)
+	assert.True(t, c2.resolved)
+}
diff --git a/internal/account_client_test.go b/internal/account_client_test.go
@@ -0,0 +1,23 @@
+package internal
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMwsAccAccountClient_GetWorkspaceClient_NoTranspile(t *testing.T) {
+	ctx, a := accountTest(t)
+	wss, err := a.Workspaces.List(ctx)
+	require.NoError(t, err)
+
+	if len(wss) == 0 {
+		t.Skip("No workspaces found")
+	}
+	w, err := a.GetWorkspaceClient(wss[0])
+	assert.NoError(t, err)
+	me, err := w.CurrentUser.Me(ctx)
+	assert.NoError(t, err)
+	assert.True(t, me.Active)
+}
diff --git a/service/provisioning/azure.go b/service/provisioning/azure.go
@@ -0,0 +1,13 @@
+package provisioning
+
+import "fmt"
+
+const resourceIdFormat = "/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Databricks/workspaces/%s"
+
+// Return the AzureResourceID for the workspace, if it is an Azure workspace.
+func (w Workspace) AzureResourceId() string {
+	if w.AzureWorkspaceInfo == nil {
+		return ""
+	}
+	return fmt.Sprintf(resourceIdFormat, w.AzureWorkspaceInfo.SubscriptionId, w.AzureWorkspaceInfo.ResourceGroup, w.WorkspaceName)
+}
diff --git a/service/provisioning/azure_test.go b/service/provisioning/azure_test.go
@@ -0,0 +1,25 @@
+package provisioning
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAzureResourceId_AzureWorkspace(t *testing.T) {
+	w := Workspace{
+		WorkspaceName: "test",
+		AzureWorkspaceInfo: &AzureWorkspaceInfo{
+			SubscriptionId: "sub",
+			ResourceGroup:  "rg",
+		},
+	}
+	assert.Equal(t, "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Databricks/workspaces/test", w.AzureResourceId())
+}
+
+func TestAzureResourceId_NonAzureWorkspace(t *testing.T) {
+	w := Workspace{
+		WorkspaceName: "test",
+	}
+	assert.Equal(t, "", w.AzureResourceId())
+}