From ae9d6883ee30182d5eeee21e30329e086243b229 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Mon, 20 Feb 2023 21:56:31 +0100
Subject: [PATCH] Complete argument for the environment flag (#221)

Command completion can be configured through `bricks completion`.
---
 cmd/root/bundle.go | 44 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/cmd/root/bundle.go b/cmd/root/bundle.go
index 40951f98..f0f8c6ad 100644
--- a/cmd/root/bundle.go
+++ b/cmd/root/bundle.go
@@ -6,6 +6,7 @@ import (
 	"github.com/databricks/bricks/bundle"
 	"github.com/databricks/bricks/bundle/config/mutator"
 	"github.com/spf13/cobra"
+	"golang.org/x/exp/maps"
 )
 
 const envName = "DATABRICKS_BUNDLE_ENV"
@@ -25,9 +26,30 @@ func getEnvironment(cmd *cobra.Command) (value string) {
 	return os.Getenv(envName)
 }
 
+// loadBundle loads the bundle configuration and applies default mutators.
+func loadBundle(cmd *cobra.Command, args []string, load func() (*bundle.Bundle, error)) (*bundle.Bundle, error) {
+	b, err := load()
+	if err != nil {
+		return nil, err
+	}
+
+	// No bundle is fine in case of `TryConfigureBundle`.
+	if b == nil {
+		return nil, nil
+	}
+
+	ctx := cmd.Context()
+	err = bundle.Apply(ctx, b, mutator.DefaultMutators())
+	if err != nil {
+		return nil, err
+	}
+
+	return b, nil
+}
+
 // configureBundle loads the bundle configuration and configures it on the command's context.
 func configureBundle(cmd *cobra.Command, args []string, load func() (*bundle.Bundle, error)) error {
-	b, err := load()
+	b, err := loadBundle(cmd, args, load)
 	if err != nil {
 		return err
 	}
@@ -37,16 +59,16 @@ func configureBundle(cmd *cobra.Command, args []string, load func() (*bundle.Bun
 		return nil
 	}
 
-	ms := mutator.DefaultMutators()
+	var m bundle.Mutator
 	env := getEnvironment(cmd)
 	if env == "" {
-		ms = append(ms, mutator.SelectDefaultEnvironment())
+		m = mutator.SelectDefaultEnvironment()
 	} else {
-		ms = append(ms, mutator.SelectEnvironment(env))
+		m = mutator.SelectEnvironment(env)
 	}
 
 	ctx := cmd.Context()
-	err = bundle.Apply(ctx, b, ms)
+	err = bundle.Apply(ctx, b, []bundle.Mutator{m})
 	if err != nil {
 		return err
 	}
@@ -66,7 +88,19 @@ func TryConfigureBundle(cmd *cobra.Command, args []string) error {
 	return configureBundle(cmd, args, bundle.TryLoad)
 }
 
+// environmentCompletion executes to autocomplete the argument to the environment flag.
+func environmentCompletion(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	b, err := loadBundle(cmd, args, bundle.MustLoad)
+	if err != nil {
+		cobra.CompErrorln(err.Error())
+		return nil, cobra.ShellCompDirectiveError
+	}
+
+	return maps.Keys(b.Config.Environments), cobra.ShellCompDirectiveDefault
+}
+
 func init() {
 	// To operate in the context of a bundle, all commands must take an "environment" parameter.
 	RootCmd.PersistentFlags().StringP("environment", "e", "", "bundle environment to use (if applicable)")
+	RootCmd.RegisterFlagCompletionFunc("environment", environmentCompletion)
 }