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

@@ -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

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
+}

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))
+}

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 {