From 5a53b118a7d2c509e7470f0fdcb96666f740b845 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 19 Oct 2023 14:34:20 +0200
Subject: [PATCH] Skip prompt on completion hook (#888)

## Changes

The first stab at this was added in #837 but only included the
`NoPrompt` check in `MustAccountClient`. I renamed it to `SkipPrompt`
(in preparation for another option that skips bundle load) and made it
work for `MustWorkspaceClient` as well.

## Tests

Manually confirmed that the completion hook no longer prompts for a
profile (when called directly with `databricks __complete`).
---
 cmd/root/auth.go              | 20 ++------------------
 cmd/root/auth_options.go      | 26 ++++++++++++++++++++++++++
 cmd/root/auth_options_test.go | 16 ++++++++++++++++
 cmd/sync/sync.go              |  3 +--
 4 files changed, 45 insertions(+), 20 deletions(-)
 create mode 100644 cmd/root/auth_options.go
 create mode 100644 cmd/root/auth_options_test.go

diff --git a/cmd/root/auth.go b/cmd/root/auth.go
index ed91090e..1e051aef 100644
--- a/cmd/root/auth.go
+++ b/cmd/root/auth.go
@@ -92,8 +92,7 @@ func MustAccountClient(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	noPrompt, ok := cmd.Context().Value(noPromptKey).(bool)
-	allowPrompt := !hasProfileFlag && (!ok || !noPrompt)
+	allowPrompt := !hasProfileFlag && !shouldSkipPrompt(cmd.Context())
 	a, err := accountClientOrPrompt(cmd.Context(), cfg, allowPrompt)
 	if err != nil {
 		return err
@@ -103,21 +102,6 @@ func MustAccountClient(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
-type noPrompt int
-
-var noPromptKey noPrompt
-
-// NoPrompt allows to skip prompt for profile configuration in MustWorkspaceClient.
-//
-// When calling MustWorkspaceClient we want to be able to customise if to show prompt or not.
-// Since we can't change function interface, in the code we only have an access to `cmd“ object.
-// Command struct does not have any state flag which indicates that it's being called in completion mode and
-// thus the Context object seems to be the only viable option for us to configure prompt behaviour based on
-// the context it's executed from.
-func NoPrompt(ctx context.Context) context.Context {
-	return context.WithValue(ctx, noPromptKey, true)
-}
-
 // Helper function to create a workspace client or prompt once if the given configuration is not valid.
 func workspaceClientOrPrompt(ctx context.Context, cfg *config.Config, allowPrompt bool) (*databricks.WorkspaceClient, error) {
 	w, err := databricks.NewWorkspaceClient((*databricks.Config)(cfg))
@@ -174,7 +158,7 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
 		cfg = currentBundle.WorkspaceClient().Config
 	}
 
-	allowPrompt := !hasProfileFlag
+	allowPrompt := !hasProfileFlag && !shouldSkipPrompt(cmd.Context())
 	w, err := workspaceClientOrPrompt(cmd.Context(), cfg, allowPrompt)
 	if err != nil {
 		return err
diff --git a/cmd/root/auth_options.go b/cmd/root/auth_options.go
new file mode 100644
index 00000000..3961f1aa
--- /dev/null
+++ b/cmd/root/auth_options.go
@@ -0,0 +1,26 @@
+package root
+
+import (
+	"context"
+)
+
+type skipPrompt int
+
+var skipPromptKey skipPrompt
+
+// SkipPrompt allows to skip prompt for profile configuration in MustWorkspaceClient.
+//
+// When calling MustWorkspaceClient we want to be able to customise if to show prompt or not.
+// Since we can't change function interface, in the code we only have an access to `cmd` object.
+// Command struct does not have any state flag which indicates that it's being called in completion mode and
+// thus the Context object seems to be the only viable option for us to configure prompt behaviour based on
+// the context it's executed from.
+func SkipPrompt(ctx context.Context) context.Context {
+	return context.WithValue(ctx, skipPromptKey, true)
+}
+
+// shouldSkipPrompt returns whether or not [SkipPrompt] has been set on the specified context.
+func shouldSkipPrompt(ctx context.Context) bool {
+	skipPrompt, ok := ctx.Value(skipPromptKey).(bool)
+	return ok && skipPrompt
+}
diff --git a/cmd/root/auth_options_test.go b/cmd/root/auth_options_test.go
new file mode 100644
index 00000000..61c68acc
--- /dev/null
+++ b/cmd/root/auth_options_test.go
@@ -0,0 +1,16 @@
+package root
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSkipPrompt(t *testing.T) {
+	ctx := context.Background()
+	assert.False(t, shouldSkipPrompt(ctx))
+
+	ctx = SkipPrompt(ctx)
+	assert.True(t, shouldSkipPrompt(ctx))
+}
diff --git a/cmd/sync/sync.go b/cmd/sync/sync.go
index 5416b573..dffa5345 100644
--- a/cmd/sync/sync.go
+++ b/cmd/sync/sync.go
@@ -149,8 +149,7 @@ func New() *cobra.Command {
 	}
 
 	cmd.ValidArgsFunction = func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-		ctx := cmd.Context()
-		cmd.SetContext(root.NoPrompt(ctx))
+		cmd.SetContext(root.SkipPrompt(cmd.Context()))
 
 		err := root.MustWorkspaceClient(cmd, args)
 		if err != nil {