From f203731fe69b047b0d26fdd28c422fe5af8e901c Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Wed, 12 Jul 2023 12:05:51 +0200
Subject: [PATCH] Support tab completion for profiles (#572)

## Changes
Currently, `databricks --profile <TAB>` autocompletes with the shell
default behavior, listing files in the local directory. This is not a
great experience. Especially given that the suggested profile names for
accounts are so long, it can be cumbersome to type them out by hand.
This PR configures autocompletion for `--profile` to inspect the
profiles of ~/.databrickscfg.

One potential improvement is to filter the response based on whether the
command is known to be account-level or workspace-level.

## Tests
Manual test.
<img width="579" alt="Screenshot_11_07_2023__18_31"
src="https://github.com/databricks/cli/assets/1850319/d7a3acd0-2511-45ac-bd82-95567775c10a">
---
 cmd/root/auth.go               |  1 +
 libs/databrickscfg/profiles.go | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/cmd/root/auth.go b/cmd/root/auth.go
index 61068ab3..ae7f7396 100644
--- a/cmd/root/auth.go
+++ b/cmd/root/auth.go
@@ -23,6 +23,7 @@ var currentUser int
 
 func init() {
 	RootCmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile")
+	RootCmd.RegisterFlagCompletionFunc("profile", databrickscfg.ProfileCompletion)
 }
 
 func MustAccountClient(cmd *cobra.Command, args []string) error {
diff --git a/libs/databrickscfg/profiles.go b/libs/databrickscfg/profiles.go
index 60b2a89a..7892bddd 100644
--- a/libs/databrickscfg/profiles.go
+++ b/libs/databrickscfg/profiles.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/spf13/cobra"
 )
 
 // Profile holds a subset of the keys in a databrickscfg profile.
@@ -59,6 +60,10 @@ func MatchAccountProfiles(p Profile) bool {
 	return p.Host != "" && p.AccountID != ""
 }
 
+func MatchAllProfiles(p Profile) bool {
+	return true
+}
+
 const DefaultPath = "~/.databrickscfg"
 
 func LoadProfiles(path string, fn ProfileMatchFunction) (file string, profiles Profiles, err error) {
@@ -99,3 +104,11 @@ func LoadProfiles(path string, fn ProfileMatchFunction) (file string, profiles P
 
 	return
 }
+
+func ProfileCompletion(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	_, profiles, err := LoadProfiles(DefaultPath, MatchAllProfiles)
+	if err != nil {
+		return nil, cobra.ShellCompDirectiveError
+	}
+	return profiles.Names(), cobra.ShellCompDirectiveNoFileComp
+}