Add command line autocomplete to the fs commands

cmd/fs/cat.go

@@ -30,5 +30,8 @@ func newCatCommand() *cobra.Command {
 		return cmdio.Render(ctx, r)
 	}
 
+	v := newValidArgs()
+	cmd.ValidArgsFunction = v.Validate
+
 	return cmd
 }

cmd/fs/cp.go

@@ -200,5 +200,10 @@ func newCpCommand() *cobra.Command {
 		return c.cpFileToFile(sourcePath, targetPath)
 	}
 
+	v := newValidArgs()
+	// The copy command has two paths that can be completed (SOURCE_PATH & TARGET_PATH)
+	v.pathArgCount = 2
+	cmd.ValidArgsFunction = v.Validate
+
 	return cmd
 }

cmd/fs/helpers.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/filer/completer"
+	"github.com/spf13/cobra"
 )
 
 func filerForPath(ctx context.Context, fullPath string) (filer.Filer, string, error) {
@@ -46,6 +48,58 @@ 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)
+}
+
+type validArgs struct {
+	mustWorkspaceClientFunc func(cmd *cobra.Command, args []string) error
+	filerForPathFunc        func(ctx context.Context, fullPath string) (filer.Filer, string, error)
+	pathArgCount            int
+	onlyDirs                bool
+}
+
+func newValidArgs() *validArgs {
+	return &validArgs{
+		mustWorkspaceClientFunc: root.MustWorkspaceClient,
+		filerForPathFunc:        filerForPath,
+		pathArgCount:            1,
+		onlyDirs:                false,
+	}
+}
+
+func (v *validArgs) Validate(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	cmd.SetContext(root.SkipPrompt(cmd.Context()))
+
+	if len(args) >= v.pathArgCount {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
+	err := v.mustWorkspaceClientFunc(cmd, args)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	filer, toCompletePath, err := v.filerForPathFunc(cmd.Context(), toComplete)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	completer := completer.New(cmd.Context(), filer, v.onlyDirs)
+
+	// Dbfs should have a prefix and always use the "/" separator
+	isDbfsPath := isDbfsPath(toComplete)
+	if isDbfsPath {
+		completer.SetPrefix(dbfsPrefix)
+		completer.SetIsLocalPath(false)
+	}
+
+	completions, directive, err := completer.CompletePath(toCompletePath)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	return completions, directive
 }

cmd/fs/helpers_test.go

@@ -3,9 +3,13 @@ package fs
 import (
 	"context"
 	"runtime"
+	"strings"
 	"testing"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -60,3 +64,88 @@ func TestFilerForWindowsLocalPaths(t *testing.T) {
 	testWindowsFilerForPath(t, ctx, `d:\abc`)
 	testWindowsFilerForPath(t, ctx, `f:\abc\ef`)
 }
+
+func mockMustWorkspaceClientFunc(cmd *cobra.Command, args []string) error {
+	return nil
+}
+
+func setupCommand(t *testing.T) (*cobra.Command, *mocks.MockWorkspaceClient) {
+	m := mocks.NewMockWorkspaceClient(t)
+	ctx := context.Background()
+	ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
+
+	cmd := &cobra.Command{}
+	cmd.SetContext(ctx)
+
+	return cmd, m
+}
+
+func setupTest(t *testing.T) (*validArgs, *cobra.Command, *mocks.MockWorkspaceClient) {
+	cmd, m := setupCommand(t)
+
+	fakeFilerForPath := func(ctx context.Context, fullPath string) (filer.Filer, string, error) {
+		fakeFiler := filer.NewFakeFiler(map[string]filer.FakeFileInfo{
+			"dir":       {FakeName: "root", FakeDir: true},
+			"dir/dirA":  {FakeDir: true},
+			"dir/dirB":  {FakeDir: true},
+			"dir/fileA": {},
+		})
+		return fakeFiler, strings.TrimPrefix(fullPath, "dbfs:/"), nil
+	}
+
+	v := newValidArgs()
+	v.filerForPathFunc = fakeFilerForPath
+	v.mustWorkspaceClientFunc = mockMustWorkspaceClientFunc
+
+	return v, cmd, m
+}
+
+func TestGetValidArgsFunctionDbfsCompletion(t *testing.T) {
+	v, cmd, _ := setupTest(t)
+	completions, directive := v.Validate(cmd, []string{}, "dbfs:/dir/")
+	assert.Equal(t, []string{"dbfs:/dir/dirA/", "dbfs:/dir/dirB/", "dbfs:/dir/fileA"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionLocalCompletion(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip()
+	}
+
+	v, cmd, _ := setupTest(t)
+	completions, directive := v.Validate(cmd, []string{}, "dir/")
+	assert.Equal(t, []string{"dir/dirA/", "dir/dirB/", "dir/fileA", "dbfs:/"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionLocalCompletionWindows(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.Skip()
+	}
+
+	v, cmd, _ := setupTest(t)
+	completions, directive := v.Validate(cmd, []string{}, "dir/")
+	assert.Equal(t, []string{"dir\\dirA\\", "dir\\dirB\\", "dir\\fileA", "dbfs:/"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionCompletionOnlyDirs(t *testing.T) {
+	v, cmd, _ := setupTest(t)
+	v.onlyDirs = true
+	completions, directive := v.Validate(cmd, []string{}, "dbfs:/dir/")
+	assert.Equal(t, []string{"dbfs:/dir/dirA/", "dbfs:/dir/dirB/"}, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoSpace, directive)
+}
+
+func TestGetValidArgsFunctionNotCompletedArgument(t *testing.T) {
+	cmd, _ := setupCommand(t)
+
+	v := newValidArgs()
+	v.pathArgCount = 0
+	v.mustWorkspaceClientFunc = mockMustWorkspaceClientFunc
+
+	completions, directive := v.Validate(cmd, []string{}, "dbfs:/")
+
+	assert.Nil(t, completions)
+	assert.Equal(t, cobra.ShellCompDirectiveNoFileComp, directive)
+}

cmd/fs/ls.go

@@ -89,5 +89,9 @@ func newLsCommand() *cobra.Command {
 		`))
 	}
 
+	v := newValidArgs()
+	v.onlyDirs = true
+	cmd.ValidArgsFunction = v.Validate
+
 	return cmd
 }

cmd/fs/mkdir.go

@@ -28,5 +28,9 @@ func newMkdirCommand() *cobra.Command {
 		return f.Mkdir(ctx, path)
 	}
 
+	v := newValidArgs()
+	v.onlyDirs = true
+	cmd.ValidArgsFunction = v.Validate
+
 	return cmd
 }

cmd/fs/rm.go

@@ -32,5 +32,8 @@ func newRmCommand() *cobra.Command {
 		return f.Delete(ctx, path)
 	}
 
+	v := newValidArgs()
+	cmd.ValidArgsFunction = v.Validate
+
 	return cmd
 }

internal/completer_test.go

@@ -0,0 +1,27 @@
+package internal
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	_ "github.com/databricks/cli/cmd/fs"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func setupCompletionFile(t *testing.T, f filer.Filer) {
+	err := f.Write(context.Background(), "dir1/file1.txt", strings.NewReader("abc"), filer.CreateParentDirectories)
+	require.NoError(t, err)
+}
+
+func TestAccFsCompletion(t *testing.T) {
+	f, tmpDir := setupDbfsFiler(t)
+	setupCompletionFile(t, f)
+
+	stdout, _ := RequireSuccessfulRun(t, "__complete", "fs", "ls", tmpDir)
+	expectedOutput := fmt.Sprintf("%s/dir1/\n:2\n", tmpDir)
+	assert.Equal(t, expectedOutput, stdout.String())
+}

libs/filer/completer/completer.go

@@ -0,0 +1,95 @@
+package completer
+
+import (
+	"context"
+	"path"
+	"path/filepath"
+	"strings"
+
+	"github.com/databricks/cli/libs/filer"
+	"github.com/spf13/cobra"
+)
+
+type completer struct {
+	ctx context.Context
+
+	// The filer to use for completing remote or local paths.
+	filer filer.Filer
+
+	// CompletePath will only return directories when onlyDirs is true.
+	onlyDirs bool
+
+	// Prefix to prepend to completions.
+	prefix string
+
+	// Whether the path is local or remote. If the path is local we use the `filepath`
+	// package for path manipulation. Otherwise we use the `path` package.
+	isLocalPath bool
+}
+
+// General completer that takes a filer to complete remote paths when TAB-ing through a path.
+func New(ctx context.Context, filer filer.Filer, onlyDirs bool) *completer {
+	return &completer{ctx: ctx, filer: filer, onlyDirs: onlyDirs, prefix: "", isLocalPath: true}
+}
+
+func (c *completer) SetPrefix(p string) {
+	c.prefix = p
+}
+
+func (c *completer) SetIsLocalPath(i bool) {
+	c.isLocalPath = i
+}
+
+func (c *completer) CompletePath(p string) ([]string, cobra.ShellCompDirective, error) {
+	trailingSeparator := "/"
+	joinFunc := path.Join
+
+	// Use filepath functions if we are in a local path.
+	if c.isLocalPath {
+		joinFunc = filepath.Join
+		trailingSeparator = string(filepath.Separator)
+	}
+
+	// If the user is TAB-ing their way through a path and the
+	// path ends in a trailing slash, we should list nested directories.
+	// If the path is incomplete, however, then we should list adjacent
+	// directories.
+	dirPath := p
+	if !strings.HasSuffix(p, trailingSeparator) {
+		dirPath = path.Dir(p)
+	}
+
+	entries, err := c.filer.ReadDir(c.ctx, dirPath)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError, err
+	}
+
+	completions := []string{}
+	for _, entry := range entries {
+		if c.onlyDirs && !entry.IsDir() {
+			continue
+		}
+
+		// Join directory path and entry name
+		completion := joinFunc(dirPath, entry.Name())
+
+		// Prepend prefix if it has been set
+		if c.prefix != "" {
+			completion = joinFunc(c.prefix, completion)
+		}
+
+		// Add trailing separator for directories.
+		if entry.IsDir() {
+			completion += trailingSeparator
+		}
+
+		completions = append(completions, completion)
+	}
+
+	// If the path is local, we add the dbfs:/ prefix suggestion as an option
+	if c.isLocalPath {
+		completions = append(completions, "dbfs:/")
+	}
+
+	return completions, cobra.ShellCompDirectiveNoSpace, err
+}