Complete positional argument to bundle run (#220)

Command completion can be configured through `bricks completion`.

bundle/run/keys.go

@@ -31,3 +31,31 @@ func ResourceKeys(b *bundle.Bundle) (keyOnly RunnerLookup, keyWithType RunnerLoo
 	}
 	return
 }
+
+// ResourceCompletions returns a list of keys that unambiguously reference resources in the bundle.
+func ResourceCompletions(b *bundle.Bundle) []string {
+	seen := make(map[string]bool)
+	comps := []string{}
+	keyOnly, keyWithType := ResourceKeys(b)
+
+	// First add resources that can be identified by key alone.
+	for k, v := range keyOnly {
+		// Invariant: len(v) >= 1. See [ResourceKeys].
+		if len(v) == 1 {
+			seen[v[0].Key()] = true
+			comps = append(comps, k)
+		}
+	}
+
+	// Then add resources that can only be identified by their type and key.
+	for k, v := range keyWithType {
+		// Invariant: len(v) == 1. See [ResourceKeys].
+		_, ok := seen[v[0].Key()]
+		if ok {
+			continue
+		}
+		comps = append(comps, k)
+	}
+
+	return comps
+}

bundle/run/keys_test.go

@@ -0,0 +1,25 @@
+package run
+
+import (
+	"testing"
+
+	"github.com/databricks/bricks/bundle"
+	"github.com/databricks/bricks/bundle/config"
+	"github.com/databricks/bricks/bundle/config/resources"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestResourceCompletionsUnique(t *testing.T) {
+	b := &bundle.Bundle{
+		Config: config.Root{
+			Resources: config.Resources{
+				Jobs: map[string]*resources.Job{
+					"foo": {},
+					"bar": {},
+				},
+			},
+		},
+	}
+
+	assert.ElementsMatch(t, []string{"foo", "bar"}, ResourceCompletions(b))
+}

cmd/bundle/run.go

@@ -40,6 +40,27 @@ var runCmd = &cobra.Command{
 
 		return nil
 	},
+
+	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		if len(args) > 0 {
+			return nil, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		err := root.MustConfigureBundle(cmd, args)
+		if err != nil {
+			cobra.CompErrorln(err.Error())
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		// No completion in the context of a bundle.
+		// Source and destination paths are taken from bundle configuration.
+		b := bundle.GetOrNil(cmd.Context())
+		if b == nil {
+			return nil, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		return run.ResourceCompletions(b), cobra.ShellCompDirectiveNoFileComp
+	},
 }
 
 func init() {