From 7139487c2f500c2a1c1ba1c8b9d4195c6633eb9d Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Thu, 19 Oct 2023 14:50:46 +0200
Subject: [PATCH] Never load authentication configuration from bundle for sync
 command (#889)

## Changes

This is used for the sync command, where we need to ensure that a bundle
configuration never taints the authentication setup as prepared in the
environment (by our VS Code extension). Once the VS Code extension fully
builds on bundles, we can remove this check again.

## Tests

Manually confirmed that calling `databricks sync` from a bundle
directory no longer picks up its authentication configuration.
---
 cmd/root/auth.go              | 19 +++++++++----------
 cmd/root/auth_options.go      | 19 +++++++++++++++++++
 cmd/root/auth_options_test.go |  8 ++++++++
 cmd/sync/sync.go              | 10 ++++++++--
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/cmd/root/auth.go b/cmd/root/auth.go
index 1e051aef..81c71479 100644
--- a/cmd/root/auth.go
+++ b/cmd/root/auth.go
@@ -146,16 +146,15 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
 		cfg.Profile = profile
 	}
 
-	// try configuring a bundle
-	err := TryConfigureBundle(cmd, args)
-	if err != nil {
-		return err
-	}
-
-	// and load the config from there
-	currentBundle := bundle.GetOrNil(cmd.Context())
-	if currentBundle != nil {
-		cfg = currentBundle.WorkspaceClient().Config
+	// Try to load a bundle configuration if we're allowed to by the caller (see `./auth_options.go`).
+	if !shouldSkipLoadBundle(cmd.Context()) {
+		err := TryConfigureBundle(cmd, args)
+		if err != nil {
+			return err
+		}
+		if b := bundle.GetOrNil(cmd.Context()); b != nil {
+			cfg = b.WorkspaceClient().Config
+		}
 	}
 
 	allowPrompt := !hasProfileFlag && !shouldSkipPrompt(cmd.Context())
diff --git a/cmd/root/auth_options.go b/cmd/root/auth_options.go
index 3961f1aa..701900d4 100644
--- a/cmd/root/auth_options.go
+++ b/cmd/root/auth_options.go
@@ -24,3 +24,22 @@ func shouldSkipPrompt(ctx context.Context) bool {
 	skipPrompt, ok := ctx.Value(skipPromptKey).(bool)
 	return ok && skipPrompt
 }
+
+type skipLoadBundle int
+
+var skipLoadBundleKey skipLoadBundle
+
+// SkipLoadBundle instructs [MustWorkspaceClient] to never try and load a bundle for configuration options.
+//
+// This is used for the sync command, where we need to ensure that a bundle configuration never taints
+// the authentication setup as prepared in the environment (by our VS Code extension).
+// Once the VS Code extension fully builds on bundles, we can remove this check again.
+func SkipLoadBundle(ctx context.Context) context.Context {
+	return context.WithValue(ctx, skipLoadBundleKey, true)
+}
+
+// shouldSkipLoadBundle returns whether or not [SkipLoadBundle] has been set on the specified context.
+func shouldSkipLoadBundle(ctx context.Context) bool {
+	skipLoadBundle, ok := ctx.Value(skipLoadBundleKey).(bool)
+	return ok && skipLoadBundle
+}
diff --git a/cmd/root/auth_options_test.go b/cmd/root/auth_options_test.go
index 61c68acc..477c2029 100644
--- a/cmd/root/auth_options_test.go
+++ b/cmd/root/auth_options_test.go
@@ -14,3 +14,11 @@ func TestSkipPrompt(t *testing.T) {
 	ctx = SkipPrompt(ctx)
 	assert.True(t, shouldSkipPrompt(ctx))
 }
+
+func TestSkipLoadBundle(t *testing.T) {
+	ctx := context.Background()
+	assert.False(t, shouldSkipLoadBundle(ctx))
+
+	ctx = SkipLoadBundle(ctx)
+	assert.True(t, shouldSkipLoadBundle(ctx))
+}
diff --git a/cmd/sync/sync.go b/cmd/sync/sync.go
index dffa5345..2870e1e0 100644
--- a/cmd/sync/sync.go
+++ b/cmd/sync/sync.go
@@ -89,7 +89,13 @@ func New() *cobra.Command {
 	cmd.Flags().BoolVar(&f.watch, "watch", false, "watch local file system for changes")
 	cmd.Flags().Var(&f.output, "output", "type of output format")
 
-	cmd.PreRunE = root.MustWorkspaceClient
+	// Wrapper for [root.MustWorkspaceClient] that disables loading authentication configuration from a bundle.
+	mustWorkspaceClient := func(cmd *cobra.Command, args []string) error {
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
+	cmd.PreRunE = mustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		var opts *sync.SyncOptions
 		var err error
@@ -151,7 +157,7 @@ func New() *cobra.Command {
 	cmd.ValidArgsFunction = func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 		cmd.SetContext(root.SkipPrompt(cmd.Context()))
 
-		err := root.MustWorkspaceClient(cmd, args)
+		err := mustWorkspaceClient(cmd, args)
 		if err != nil {
 			return nil, cobra.ShellCompDirectiveError
 		}