From dd95668474ab4369eaebe40f78df5146c4fd5970 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Mon, 20 Feb 2023 21:55:06 +0100
Subject: [PATCH] Complete positional argument to bundle run (#220)

Command completion can be configured through `bricks completion`.
---
 bundle/run/keys.go      | 28 ++++++++++++++++++++++++++++
 bundle/run/keys_test.go | 25 +++++++++++++++++++++++++
 cmd/bundle/run.go       | 21 +++++++++++++++++++++
 3 files changed, 74 insertions(+)
 create mode 100644 bundle/run/keys_test.go

diff --git a/bundle/run/keys.go b/bundle/run/keys.go
index 581b6587..b7d152b9 100644
--- a/bundle/run/keys.go
+++ b/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
+}
diff --git a/bundle/run/keys_test.go b/bundle/run/keys_test.go
new file mode 100644
index 00000000..9cbd033b
--- /dev/null
+++ b/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))
+}
diff --git a/cmd/bundle/run.go b/cmd/bundle/run.go
index 335c6954..cc358c7c 100644
--- a/cmd/bundle/run.go
+++ 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() {