Add command line autocomplete to the fs commands

cmd/fs/cat.go

@@ -30,5 +30,7 @@ func newCatCommand() *cobra.Command {
 		return cmdio.Render(ctx, r)
 	}
 
+	cmd.ValidArgsFunction = getValidArgsFunction(1, false, filerForPath)
+
 	return cmd
 }

cmd/fs/cp.go

@@ -200,5 +200,8 @@ func newCpCommand() *cobra.Command {
 		return c.cpFileToFile(sourcePath, targetPath)
 	}
 
+	// The copy command has two paths that can be completed (SOURCE_PATH & TARGET_PATH)
+	cmd.ValidArgsFunction = getValidArgsFunction(2, false, filerForPath)
+
 	return cmd
 }

cmd/fs/helpers.go

@@ -7,7 +7,9 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/completer"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/spf13/cobra"
 )
 
 func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, error) {
@@ -46,6 +48,65 @@ func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, er
 	return f, path, err
 }
 
+const dbfsPrefix string = "dbfs:/"
+
 func isDbfsPath(path string) bool {
-	return strings.HasPrefix(path, "dbfs:/")
+	return strings.HasPrefix(path, dbfsPrefix)
+}
+
+func getValidArgsFunction(pathArgCount int, onlyDirs bool, filerForPathFunc func(ctx context.Context, fullPath string) (filer.Filer, string, error)) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		cmd.SetContext(root.SkipPrompt(cmd.Context()))
+
+		err := mustWorkspaceClient(cmd, args) // TODO: Fix this
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		isValidPrefix := isDbfsPath(toComplete)
+
+		if !isValidPrefix {
+			return []string{dbfsPrefix}, cobra.ShellCompDirectiveNoSpace
+		}
+
+		filer, path, err := filerForPathFunc(cmd.Context(), toComplete)
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		wsc := root.WorkspaceClient(cmd.Context())
+		_, err = wsc.CurrentUser.Me(cmd.Context())
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		completer := completer.NewCompleter(cmd.Context(), filer, onlyDirs)
+
+		if len(args) < pathArgCount {
+			completions, directive := completer.CompleteRemotePath(path)
+
+			// The completions will start with a "/", so we'll prefix
+			// them with the dbfsPrefix without a trailing "/" (dbfs:)
+			prefix := dbfsPrefix[:len(dbfsPrefix)-1]
+
+			// Add the prefix to the completions
+			for i, completion := range completions {
+				completions[i] = fmt.Sprintf("%s%s", prefix, completion)
+			}
+
+			return completions, directive
+		} else {
+			return nil, cobra.ShellCompDirectiveNoFileComp
+		}
+	}
+}
+
+// Wrapper for [root.MustWorkspaceClient] that disables loading authentication configuration from a bundle.
+func mustWorkspaceClient(cmd *cobra.Command, args []string) error {
+	if root.HasWorkspaceClient(cmd.Context()) {
+		return nil
+	}
+
+	cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+	return root.MustWorkspaceClient(cmd, args)
 }

cmd/fs/helpers_test.go

@@ -2,11 +2,19 @@ package fs
 
 import (
 	"context"
+	"errors"
+	"io/fs"
 	"runtime"
 	"testing"
 
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -60,3 +68,92 @@ func TestFilerForWindowsLocalPaths(t *testing.T) {
 	testWindowsFilerForPath(t, ctx, `d:\abc`)
 	testWindowsFilerForPath(t, ctx, `f:\abc\ef`)
 }
+
+func mockCurrentUserApi(m *mocks.MockWorkspaceClient, u *iam.User, e error) {
+	currentUserApi := m.GetMockCurrentUserAPI()
+	currentUserApi.EXPECT().Me(mock.AnythingOfType("*context.valueCtx")).Return(u, e)
+}
+
+func setupValidArgsFunctionTest(t *testing.T) (*mocks.MockWorkspaceClient, *cobra.Command) {
+	m := mocks.NewMockWorkspaceClient(t)
+	ctx := context.Background()
+	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+
+	cmd := &cobra.Command{}
+	cmd.SetContext(ctx)
+
+	return m, cmd
+}
+
+func TestGetValidArgsFunctionCompletion(t *testing.T) {
+	m, cmd := setupValidArgsFunctionTest(t)
+
+	mockCurrentUserApi(m, nil, nil)
+
+	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+		testutil.NewFakeDirEntry("dir", true),
+		testutil.NewFakeDirEntry("file", false),
+	})
+
+	validArgsFunction := getValidArgsFunction(1, false, mockFilerForPath)
+	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
+
+	assert.Equal(t, []string{"dbfs:/dir", "dbfs:/file"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionCompletionOnlyDirs(t *testing.T) {
+	m, cmd := setupValidArgsFunctionTest(t)
+
+	mockCurrentUserApi(m, nil, nil)
+
+	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+		testutil.NewFakeDirEntry("dir", true),
+		testutil.NewFakeDirEntry("file", false),
+	})
+
+	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath)
+	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
+
+	assert.Equal(t, []string{"dbfs:/dir"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionNoDbfsPath(t *testing.T) {
+	_, cmd := setupValidArgsFunctionTest(t)
+
+	mockFilerForPath := testutil.GetMockFilerForPath(t, nil)
+
+	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath)
+	completions, directive := validArgsFunction(cmd, []string{}, "/")
+
+	assert.Equal(t, []string{"dbfs:/"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionNoCurrentUser(t *testing.T) {
+	m, cmd := setupValidArgsFunctionTest(t)
+
+	mockCurrentUserApi(m, nil, errors.New("Current User Error"))
+	mockFilerForPath := testutil.GetMockFilerForPath(t, nil)
+
+	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath)
+	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
+
+	assert.Nil(t, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveError, directive)
+}
+
+func TestGetValidArgsFunctionNotCompletedArgument(t *testing.T) {
+	m, cmd := setupValidArgsFunctionTest(t)
+
+	mockCurrentUserApi(m, nil, nil)
+	mockFilerForPath := testutil.GetMockFilerForPath(t, nil)
+
+	// 0 here means we don't complete any arguments
+	validArgsFunction := getValidArgsFunction(0, true, mockFilerForPath)
+	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
+
+	assert.Nil(t, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoFileComp, directive)
+}

cmd/fs/ls.go

@@ -89,5 +89,7 @@ func newLsCommand() *cobra.Command {
 		`))
 	}
 
+	cmd.ValidArgsFunction = getValidArgsFunction(1, true, filerForPath)
+
 	return cmd
 }

cmd/fs/mkdir.go

@@ -28,5 +28,7 @@ func newMkdirCommand() *cobra.Command {
 		return f.Mkdir(ctx, path)
 	}
 
+	cmd.ValidArgsFunction = getValidArgsFunction(1, true, filerForPath)
+
 	return cmd
 }

cmd/fs/rm.go

@@ -32,5 +32,7 @@ func newRmCommand() *cobra.Command {
 		return f.Delete(ctx, path)
 	}
 
+	cmd.ValidArgsFunction = getValidArgsFunction(1, false, filerForPath)
+
 	return cmd
 }

cmd/root/auth.go

@@ -152,6 +152,7 @@ func MustAccountClient(cmd *cobra.Command, args []string) error {
 // Helper function to create a workspace client or prompt once if the given configuration is not valid.
 func workspaceClientOrPrompt(ctx context.Context, cfg *config.Config, allowPrompt bool) (*databricks.WorkspaceClient, error) {
 	w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg))
+
 	if err == nil {
 		err = w.Config.Authenticate(emptyHttpRequest(ctx))
 	}
@@ -230,6 +231,11 @@ func SetWorkspaceClient(ctx context.Context, w *databricks.WorkspaceClient) cont
 	return context.WithValue(ctx, &workspaceClient, w)
 }
 
+func HasWorkspaceClient(ctx context.Context) bool {
+	_, ok := ctx.Value(&workspaceClient).(*databricks.WorkspaceClient)
+	return ok
+}
+
 func SetAccountClient(ctx context.Context, a *databricks.AccountClient) context.Context {
 	return context.WithValue(ctx, &accountClient, a)
 }

cmd/root/auth_options.go

@@ -6,7 +6,7 @@ import (
 
 type skipPrompt int
 
-var skipPromptKey skipPrompt
+var SkipPromptKey skipPrompt
 
 // SkipPrompt allows to skip prompt for profile configuration in MustWorkspaceClient.
 //
@@ -16,12 +16,12 @@ var skipPromptKey skipPrompt
 // thus the Context object seems to be the only viable option for us to configure prompt behaviour based on
 // the context it's executed from.
 func SkipPrompt(ctx context.Context) context.Context {
-	return context.WithValue(ctx, skipPromptKey, true)
+	return context.WithValue(ctx, SkipPromptKey, true)
 }
 
 // shouldSkipPrompt returns whether or not [SkipPrompt] has been set on the specified context.
 func shouldSkipPrompt(ctx context.Context) bool {
-	skipPrompt, ok := ctx.Value(skipPromptKey).(bool)
+	skipPrompt, ok := ctx.Value(SkipPromptKey).(bool)
 	return ok && skipPrompt
 }
 

internal/testutil/filer.go

@@ -0,0 +1,74 @@
+package testutil
+
+import (
+	"context"
+	"io/fs"
+	"testing"
+	"time"
+
+	mockfiler "github.com/databricks/cli/internal/mocks/libs/filer"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/stretchr/testify/mock"
+)
+
+type fakeDirEntry struct {
+	fakeFileInfo
+}
+
+func (entry fakeDirEntry) Type() fs.FileMode {
+	typ := fs.ModePerm
+	if entry.dir {
+		typ |= fs.ModeDir
+	}
+	return typ
+}
+
+func (entry fakeDirEntry) Info() (fs.FileInfo, error) {
+	return entry.fakeFileInfo, nil
+}
+
+type fakeFileInfo struct {
+	name string
+	size int64
+	dir  bool
+	mode fs.FileMode
+}
+
+func (info fakeFileInfo) Name() string {
+	return info.name
+}
+
+func (info fakeFileInfo) Size() int64 {
+	return info.size
+}
+
+func (info fakeFileInfo) Mode() fs.FileMode {
+	return info.mode
+}
+
+func (info fakeFileInfo) ModTime() time.Time {
+	return time.Now()
+}
+
+func (info fakeFileInfo) IsDir() bool {
+	return info.dir
+}
+
+func (info fakeFileInfo) Sys() any {
+	return nil
+}
+
+func NewFakeDirEntry(name string, dir bool) fs.DirEntry {
+	return fakeDirEntry{fakeFileInfo{name: name, dir: dir}}
+}
+
+func GetMockFilerForPath(t *testing.T, entries []fs.DirEntry) func(ctx context.Context, fullPath string) (filer.Filer, string, error) {
+	mockFiler := mockfiler.NewMockFiler(t)
+	if entries != nil {
+		mockFiler.EXPECT().ReadDir(mock.AnythingOfType("*context.valueCtx"), mock.Anything).Return(entries, nil)
+	}
+	mockFilerForPath := func(ctx context.Context, fullPath string) (filer.Filer, string, error) {
+		return mockFiler, "/", nil
+	}
+	return mockFilerForPath
+}