Move WorkspaceClient to libs/command (#2444)

## Changes
Move the WorkspaceClient reader and setter from the root package to
`libs/command`.

## Why
This allows us to read the workspace client set in the context in all
CLI packages. This was not possible before because using
`root.WorkspaceClient` would often result in an import cycle.

This would also allow us to standardise reads of the workspace client to
be from the context in bundle commands. Today those are read from the
bundle object tree instead.

This is a natural followup to
#2440.

## Tests
Existing tests and one new unit test.

NO_CHANGELOG=true

.codegen/service.go.tmpl

@@ -4,6 +4,7 @@ package {{(.TrimPrefix "account").SnakeName}}
 
 import (
     "github.com/databricks/cli/libs/cmdio"
+    "github.com/databricks/cli/libs/command"
     "github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/diag"
 	"github.com/databricks/cli/cmd/root"
@@ -240,7 +241,7 @@ func new{{.PascalName}}() *cobra.Command {
 	cmd.PreRunE = root.Must{{if .Service.IsAccounts}}Account{{else}}Workspace{{end}}Client
 	cmd.RunE = func(cmd *cobra.Command, args []string) (err error) {
 		ctx := cmd.Context()
-		{{if .Service.IsAccounts}}a := root.AccountClient(ctx){{else}}w := root.WorkspaceClient(ctx){{end}}
+		{{if .Service.IsAccounts}}a := root.AccountClient(ctx){{else}}w := command.WorkspaceClient(ctx){{end}}
 		{{- if .Request }}
 			{{ if $canUseJson }}
 			if cmd.Flags().Changed("json") {

cmd/auth/describe.go

@@ -110,7 +110,7 @@ func getAuthStatus(cmd *cobra.Command, args []string, showSensitive bool, fn try
 		return &status, nil
 	}
 
-	w := root.WorkspaceClient(ctx)
+	w := command.WorkspaceClient(ctx)
 	me, err := w.CurrentUser.Me(ctx)
 	if err != nil {
 		return &authStatus{

cmd/auth/describe_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/iam"
@@ -17,7 +18,7 @@ import (
 func TestGetWorkspaceAuthStatus(t *testing.T) {
 	ctx := context.Background()
 	m := mocks.NewMockWorkspaceClient(t)
-	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+	ctx = command.SetWorkspaceClient(ctx, m.WorkspaceClient)
 
 	cmd := &cobra.Command{}
 	cmd.SetContext(ctx)
@@ -75,7 +76,7 @@ func TestGetWorkspaceAuthStatus(t *testing.T) {
 func TestGetWorkspaceAuthStatusError(t *testing.T) {
 	ctx := context.Background()
 	m := mocks.NewMockWorkspaceClient(t)
-	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+	ctx = command.SetWorkspaceClient(ctx, m.WorkspaceClient)
 
 	cmd := &cobra.Command{}
 	cmd.SetContext(ctx)
@@ -124,7 +125,7 @@ func TestGetWorkspaceAuthStatusError(t *testing.T) {
 func TestGetWorkspaceAuthStatusSensitive(t *testing.T) {
 	ctx := context.Background()
 	m := mocks.NewMockWorkspaceClient(t)
-	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+	ctx = command.SetWorkspaceClient(ctx, m.WorkspaceClient)
 
 	cmd := &cobra.Command{}
 	cmd.SetContext(ctx)

cmd/fs/helpers.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/filer/completer"
 	"github.com/spf13/cobra"
@@ -35,7 +36,7 @@ func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, er
 	}
 
 	path := parts[1]
-	w := root.WorkspaceClient(ctx)
+	w := command.WorkspaceClient(ctx)
 
 	// If the specified path has the "Volumes" prefix, use the Files API.
 	if strings.HasPrefix(path, "/Volumes/") {

cmd/fs/helpers_test.go

@@ -6,7 +6,7 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/cli/libs/fakefs"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
@@ -73,7 +73,7 @@ func mockMustWorkspaceClientFunc(cmd *cobra.Command, args []string) error {
 func setupCommand(t *testing.T) (*cobra.Command, *mocks.MockWorkspaceClient) {
 	m := mocks.NewMockWorkspaceClient(t)
 	ctx := context.Background()
-	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+	ctx = command.SetWorkspaceClient(ctx, m.WorkspaceClient)
 
 	cmd := &cobra.Command{}
 	cmd.SetContext(ctx)

cmd/root/auth.go

@@ -17,8 +17,7 @@ import (
 
 // Placeholders to use as unique keys in context.Context.
 var (
-	workspaceClient int
-	accountClient   int
+	accountClient int
 )
 
 type ErrNoWorkspaceProfiles struct {
@@ -229,15 +228,11 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	ctx = context.WithValue(ctx, &workspaceClient, w)
+	ctx = command.SetWorkspaceClient(ctx, w)
 	cmd.SetContext(ctx)
 	return nil
 }
 
-func SetWorkspaceClient(ctx context.Context, w *databricks.WorkspaceClient) context.Context {
-	return context.WithValue(ctx, &workspaceClient, w)
-}
-
 func SetAccountClient(ctx context.Context, a *databricks.AccountClient) context.Context {
 	return context.WithValue(ctx, &accountClient, a)
 }
@@ -321,14 +316,6 @@ func emptyHttpRequest(ctx context.Context) *http.Request {
 	return req
 }
 
-func WorkspaceClient(ctx context.Context) *databricks.WorkspaceClient {
-	w, ok := ctx.Value(&workspaceClient).(*databricks.WorkspaceClient)
-	if !ok {
-		panic("cannot get *databricks.WorkspaceClient. Please report it as a bug")
-	}
-	return w
-}
-
 func AccountClient(ctx context.Context) *databricks.AccountClient {
 	a, ok := ctx.Value(&accountClient).(*databricks.AccountClient)
 	if !ok {

cmd/root/auth_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -263,7 +264,7 @@ func TestMustAnyClientCanCreateWorkspaceClient(t *testing.T) {
 	require.False(t, isAccount)
 	require.NoError(t, err)
 
-	w := WorkspaceClient(cmd.Context())
+	w := command.WorkspaceClient(cmd.Context())
 	require.NotNil(t, w)
 }
 

cmd/sync/sync.go

@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/deploy/files"
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/git"
 	"github.com/databricks/cli/libs/log"
@@ -65,7 +66,7 @@ func (f *syncFlags) syncOptionsFromArgs(cmd *cobra.Command, args []string) (*syn
 	}
 
 	ctx := cmd.Context()
-	client := root.WorkspaceClient(ctx)
+	client := command.WorkspaceClient(ctx)
 
 	localRoot := vfs.MustNew(args[0])
 	info, err := git.FetchRepositoryInfo(ctx, localRoot.Native(), client)
@@ -186,7 +187,7 @@ func New() *cobra.Command {
 		case 0:
 			return nil, cobra.ShellCompDirectiveFilterDirs
 		case 1:
-			wsc := root.WorkspaceClient(cmd.Context())
+			wsc := command.WorkspaceClient(cmd.Context())
 			return completeRemotePath(cmd.Context(), wsc, toComplete)
 		default:
 			return nil, cobra.ShellCompDirectiveNoFileComp

cmd/sync/sync_test.go

@@ -8,7 +8,7 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
-	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/command"
 	"github.com/databricks/cli/libs/vfs"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -58,7 +58,7 @@ func TestSyncOptionsFromArgs(t *testing.T) {
 
 	f := syncFlags{}
 	cmd := New()
-	cmd.SetContext(root.SetWorkspaceClient(context.Background(), nil))
+	cmd.SetContext(command.SetWorkspaceClient(context.Background(), nil))
 	opts, err := f.syncOptionsFromArgs(cmd, []string{local, remote})
 	require.NoError(t, err)
 	assert.Equal(t, local, opts.LocalRoot.Native())