Add command line autocomplete to the fs commands

cmd/fs/cat.go

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

cmd/fs/cp.go

@@ -201,7 +201,7 @@ func newCpCommand() *cobra.Command {
 	}
 
 	// The copy command has two paths that can be completed (SOURCE_PATH & TARGET_PATH)
-	cmd.ValidArgsFunction = getValidArgsFunction(2, false, filerForPath)
+	cmd.ValidArgsFunction = getValidArgsFunction(2, false, filerForPath, root.MustWorkspaceClient)
 
 	return cmd
 }

cmd/fs/helpers.go

@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/completer"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 )
 
@@ -50,41 +51,34 @@ func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, er
 
 const dbfsPrefix string = "dbfs:/"
 const volumesPefix string = "dbfs:/Volumes"
+const localPefix string = "./"
 
 func isDbfsPath(path string) bool {
 	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) {
+func getValidArgsFunction(
+	pathArgCount int,
+	onlyDirs bool,
+	filerForPathFunc func(ctx context.Context, fullPath string) (filer.Filer, string, error),
+	mustWorkspaceClientFunc func(cmd *cobra.Command, args []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
+		err := mustWorkspaceClientFunc(cmd, args)
+
 		if err != nil {
 			return nil, cobra.ShellCompDirectiveError
 		}
 
-		prefixes := []string{
-			dbfsPrefix,
-			"/",
-		}
-
-		selectedPrefix := ""
-		for _, p := range prefixes {
-			if strings.HasPrefix(toComplete, p) {
-				selectedPrefix = p
-			}
-		}
-
-		if selectedPrefix == "" {
-			return prefixes, cobra.ShellCompDirectiveNoSpace
-		}
-
 		filer, path, err := filerForPathFunc(cmd.Context(), toComplete)
 		if err != nil {
 			return nil, cobra.ShellCompDirectiveError
 		}
 
+		isDbfsPath := isDbfsPath(toComplete)
+
 		wsc := root.WorkspaceClient(cmd.Context())
 		_, err = wsc.CurrentUser.Me(cmd.Context())
 		if err != nil {
@@ -99,30 +93,42 @@ func getValidArgsFunction(pathArgCount int, onlyDirs bool, filerForPathFunc func
 
 		completions, directive := completer.CompleteRemotePath(path)
 
-		// The completions will start with a "/", so we'll prefix
-		// them with the selectedPrefix without a trailing "/"
-		prefix := selectedPrefix[:len(selectedPrefix)-1]
-
-		// Add the prefix to the completions
 		for i, completion := range completions {
-			completions[i] = fmt.Sprintf("%s%s", prefix, completion)
+			if isDbfsPath {
+				// The completions will start with a "/", so we'll prefix them with the
+				// selectedPrefix without a trailing "/"
+				prefix := dbfsPrefix[:len(dbfsPrefix)-1]
+				completions[i] = fmt.Sprintf("%s%s", prefix, completion)
+			} else if shouldDropLocalPrefix(toComplete, completion) {
+
+				completions[i] = completion[len(localPefix):]
+			} else {
+				completions[i] = completion
+			}
+		}
+
+		// If the path is a dbfs path, we try to add the dbfs:/Volumes prefix suggestion
+		if isDbfsPath && strings.HasPrefix(volumesPefix, toComplete) {
+			completions = append([]string{highlight(volumesPefix)}, completions...)
+
 		}
 
-		// If the path is a dbfs path, we should also add the Volumes prefix
-		if strings.HasPrefix(volumesPefix, toComplete) {
-			completions = append(completions, volumesPefix)
+		// If the path is local, we try to prepend the dbfs:/ prefix suggestion
+		if !isDbfsPath && strings.HasPrefix(dbfsPrefix, toComplete) {
+			completions = append([]string{highlight(dbfsPrefix)}, completions...)
 		}
 
 		return completions, directive
 	}
 }
 
-// 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
-	}
+// Drop the local prefix from completions if the path to complete doesn't
+// start with it. We do this because the local filer returns paths with the
+// local prefix.
+func shouldDropLocalPrefix(toComplete string, completion string) bool {
+	return !strings.HasPrefix(toComplete, localPefix) && strings.HasPrefix(completion, localPefix)
+}
 
-	cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
-	return root.MustWorkspaceClient(cmd, args)
+func highlight(c string) string {
+	return color.New(color.FgCyan, color.Bold).Sprintf(c)
 }

cmd/fs/helpers_test.go

@@ -74,6 +74,10 @@ func mockCurrentUserApi(m *mocks.MockWorkspaceClient, u *iam.User, e error) {
 	currentUserApi.EXPECT().Me(mock.AnythingOfType("*context.valueCtx")).Return(u, e)
 }
 
+func mockMustWorkspaceClientFunc(cmd *cobra.Command, args []string) error {
+	return nil
+}
+
 func setupValidArgsFunctionTest(t *testing.T) (*mocks.MockWorkspaceClient, *cobra.Command) {
 	m := mocks.NewMockWorkspaceClient(t)
 	ctx := context.Background()
@@ -90,16 +94,16 @@ func TestGetValidArgsFunctionCompletion(t *testing.T) {
 
 	mockCurrentUserApi(m, nil, nil)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/", []fs.DirEntry{
 		testutil.NewFakeDirEntry("dir", true),
 		testutil.NewFakeDirEntry("file", false),
 	})
 
-	validArgsFunction := getValidArgsFunction(1, false, mockFilerForPath)
+	validArgsFunction := getValidArgsFunction(1, false, mockFilerForPath, mockMustWorkspaceClientFunc)
 	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
 
 	// dbfs:/Volumes is programmatically added to the completions
-	assert.Equal(t, []string{"dbfs:/dir", "dbfs:/file", "dbfs:/Volumes"}, completions)
+	assert.Equal(t, []string{"dbfs:/Volumes", "dbfs:/dir", "dbfs:/file"}, completions)
 	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
 }
 
@@ -108,28 +112,33 @@ func TestGetValidArgsFunctionCompletionOnlyDirs(t *testing.T) {
 
 	mockCurrentUserApi(m, nil, nil)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/", []fs.DirEntry{
 		testutil.NewFakeDirEntry("dir", true),
 		testutil.NewFakeDirEntry("file", false),
 	})
 
-	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath)
+	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath, mockMustWorkspaceClientFunc)
 	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
 
-	assert.Equal(t, []string{"dbfs:/dir", "dbfs:/Volumes"}, completions)
+	assert.Equal(t, []string{"dbfs:/Volumes", "dbfs:/dir"}, completions)
 	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
 }
 
 func TestGetValidArgsFunctionNoPath(t *testing.T) {
-	_, cmd := setupValidArgsFunctionTest(t)
+	m, cmd := setupValidArgsFunctionTest(t)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, nil)
+	mockCurrentUserApi(m, nil, nil)
 
-	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath)
-	completions, directive := validArgsFunction(cmd, []string{}, "")
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "", []fs.DirEntry{
+		testutil.NewFakeDirEntry("dFile1", false),
+		testutil.NewFakeDirEntry("dFile2", false),
+	})
+
+	validArgsFunction := getValidArgsFunction(1, false, mockFilerForPath, mockMustWorkspaceClientFunc)
+	completions, directive := validArgsFunction(cmd, []string{}, "d")
 
 	// Suggest both dbfs and local paths at beginning of completion
-	assert.Equal(t, []string{"dbfs:/", "/"}, completions)
+	assert.Equal(t, []string{"dbfs:/", "dFile1", "dFile2"}, completions)
 	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
 }
 
@@ -138,25 +147,42 @@ func TestGetValidArgsFunctionLocalPath(t *testing.T) {
 
 	mockCurrentUserApi(m, nil, nil)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "", []fs.DirEntry{
+		testutil.NewFakeDirEntry("dir", true),
+		testutil.NewFakeDirEntry("file", false),
+	})
+
+	validArgsFunction := getValidArgsFunction(1, false, mockFilerForPath, mockMustWorkspaceClientFunc)
+	completions, directive := validArgsFunction(cmd, []string{}, "")
+
+	assert.Equal(t, []string{"dbfs:/", "dir", "file"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionAbsoluteLocalPath(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)
+	validArgsFunction := getValidArgsFunction(1, false, mockFilerForPath, mockMustWorkspaceClientFunc)
 	completions, directive := validArgsFunction(cmd, []string{}, "/")
 
-	assert.Equal(t, []string{"/dir", "/file"}, completions)
+	assert.Equal(t, []string{"dir", "file"}, 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)
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/", nil)
 
-	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath)
+	validArgsFunction := getValidArgsFunction(1, true, mockFilerForPath, mockMustWorkspaceClientFunc)
 	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
 
 	assert.Nil(t, completions)
@@ -167,10 +193,10 @@ func TestGetValidArgsFunctionNotCompletedArgument(t *testing.T) {
 	m, cmd := setupValidArgsFunctionTest(t)
 
 	mockCurrentUserApi(m, nil, nil)
-	mockFilerForPath := testutil.GetMockFilerForPath(t, nil)
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/", nil)
 
 	// 0 here means we don't complete any arguments
-	validArgsFunction := getValidArgsFunction(0, true, mockFilerForPath)
+	validArgsFunction := getValidArgsFunction(0, true, mockFilerForPath, mockMustWorkspaceClientFunc)
 	completions, directive := validArgsFunction(cmd, []string{}, "dbfs:/")
 
 	assert.Nil(t, completions)

cmd/fs/ls.go

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

cmd/fs/mkdir.go

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

cmd/fs/rm.go

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

internal/testutil/filer.go

@@ -62,13 +62,13 @@ 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) {
+func GetMockFilerForPath(t *testing.T, path string, 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)
+		mockFiler.EXPECT().ReadDir(mock.AnythingOfType("*context.valueCtx"), path).Return(entries, nil)
 	}
 	mockFilerForPath := func(ctx context.Context, fullPath string) (filer.Filer, string, error) {
-		return mockFiler, "/", nil
+		return mockFiler, path, nil
 	}
 	return mockFilerForPath
 }

libs/completer/completer.go

@@ -4,7 +4,7 @@ import (
 	"context"
 	"fmt"
 	"path"
-	"sort"
+	"strings"
 
 	"github.com/databricks/cli/libs/filer"
 	"github.com/spf13/cobra"
@@ -57,8 +57,8 @@ func fetchCompletions(
 
 				// Ensure the path and name have a "/" separating them. We don't use path
 				// utilities to concatenate the path and name because we want to preserve
-				// the formatting of the path the user has typed (e.g. //a/b///c)
-				if remotePath[len(remotePath)-1] != '/' {
+				// the formatting of whatever path the user has typed (e.g. //a/b///c)
+				if len(remotePath) > 0 && !strings.HasSuffix(remotePath, "/") {
 					separator = "/"
 				}
 
@@ -67,11 +67,6 @@ func fetchCompletions(
 			}
 		}
 
-		// Sort completions alphabetically
-		sort.Slice(completions, func(i, j int) bool {
-			return completions[i] < completions[j]
-		})
-
 		ch <- completions
 	}()
 

libs/completer/completer_test.go

@@ -25,7 +25,7 @@ func setup(t *testing.T) context.Context {
 func TestFilerCompleterReturnsNestedDirs(t *testing.T) {
 	ctx := setup(t)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/", []fs.DirEntry{
 		testutil.NewFakeDirEntry("dir", true),
 	})
 	mockFiler, _, _ := mockFilerForPath(ctx, "dbfs:/")
@@ -76,7 +76,7 @@ func TestFilerCompleterReadDirError(t *testing.T) {
 func TestFilerCompleterReturnsFileAndDir(t *testing.T) {
 	ctx := setup(t)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/", []fs.DirEntry{
 		testutil.NewFakeDirEntry("dir", true),
 		testutil.NewFakeDirEntry("file", false),
 	})
@@ -94,7 +94,7 @@ func TestFilerCompleterReturnsFileAndDir(t *testing.T) {
 func TestFilerCompleterRetainsFormatting(t *testing.T) {
 	ctx := setup(t)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "//dir//", []fs.DirEntry{
 		testutil.NewFakeDirEntry("nested_dir", true),
 	})
 	mockFiler, _, _ := mockFilerForPath(ctx, "dbfs://dir")
@@ -110,7 +110,7 @@ func TestFilerCompleterRetainsFormatting(t *testing.T) {
 func TestFilerCompleterAddsSeparator(t *testing.T) {
 	ctx := setup(t)
 
-	mockFilerForPath := testutil.GetMockFilerForPath(t, []fs.DirEntry{
+	mockFilerForPath := testutil.GetMockFilerForPath(t, "/dir", []fs.DirEntry{
 		testutil.NewFakeDirEntry("nested_dir", true),
 	})
 	mockFiler, _, _ := mockFilerForPath(ctx, "dbfs:/dir")