This commit is contained in:
Shreyas Goenka 2025-01-29 12:12:34 +01:00
commit ba55f5bb9c
No known key found for this signature in database
GPG Key ID: 92A07DF49CCB0622
6 changed files with 69 additions and 92 deletions

View File

@ -72,6 +72,7 @@ type Bundle struct {
// It can be initialized on demand after loading the configuration. // It can be initialized on demand after loading the configuration.
clientOnce sync.Once clientOnce sync.Once
client *databricks.WorkspaceClient client *databricks.WorkspaceClient
clientErr error
// Files that are synced to the workspace.file_path // Files that are synced to the workspace.file_path
Files []fileset.File Files []fileset.File
@ -134,23 +135,25 @@ func TryLoad(ctx context.Context) (*Bundle, error) {
return Load(ctx, root) return Load(ctx, root)
} }
func (b *Bundle) InitializeWorkspaceClient() (*databricks.WorkspaceClient, error) { func (b *Bundle) WorkspaceClientE() (*databricks.WorkspaceClient, error) {
client, err := b.Config.Workspace.Client() b.clientOnce.Do(func() {
var err error
b.client, err = b.Config.Workspace.Client()
if err != nil { if err != nil {
return nil, fmt.Errorf("cannot resolve bundle auth configuration: %w", err) b.clientErr = fmt.Errorf("cannot resolve bundle auth configuration: %w", err)
} }
return client, nil })
return b.client, b.clientErr
} }
func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient { func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
b.clientOnce.Do(func() { client, err := b.WorkspaceClientE()
var err error
b.client, err = b.InitializeWorkspaceClient()
if err != nil { if err != nil {
panic(err) panic(err)
} }
})
return b.client return client
} }
// SetWorkpaceClient sets the workspace client for this bundle. // SetWorkpaceClient sets the workspace client for this bundle.

View File

@ -1,26 +0,0 @@
package mutator
import (
"context"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/libs/diag"
)
type initializeWorkspaceClient struct{}
func InitializeWorkspaceClient() bundle.Mutator {
return &initializeWorkspaceClient{}
}
func (m *initializeWorkspaceClient) Name() string {
return "InitializeWorkspaceClient"
}
// Apply initializes the workspace client for the bundle. We do this here so
// downstream calls to b.WorkspaceClient() do not panic if there's an error in the
// auth configuration.
func (m *initializeWorkspaceClient) Apply(_ context.Context, b *bundle.Bundle) diag.Diagnostics {
_, err := b.InitializeWorkspaceClient()
return diag.FromErr(err)
}

View File

@ -34,7 +34,6 @@ func Initialize() bundle.Mutator {
// If it is an ancestor, this updates all paths to be relative to the sync root path. // If it is an ancestor, this updates all paths to be relative to the sync root path.
mutator.SyncInferRoot(), mutator.SyncInferRoot(),
mutator.InitializeWorkspaceClient(),
mutator.PopulateCurrentUser(), mutator.PopulateCurrentUser(),
mutator.LoadGitDetails(), mutator.LoadGitDetails(),

View File

@ -209,7 +209,7 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
if b != nil { if b != nil {
ctx = context.WithValue(ctx, &configUsed, b.Config.Workspace.Config()) ctx = context.WithValue(ctx, &configUsed, b.Config.Workspace.Config())
cmd.SetContext(ctx) cmd.SetContext(ctx)
client, err := b.InitializeWorkspaceClient() client, err := b.WorkspaceClientE()
if err != nil { if err != nil {
return err return err
} }

View File

@ -81,6 +81,22 @@ func configureBundle(cmd *cobra.Command, b *bundle.Bundle) (*bundle.Bundle, diag
// Configure the workspace profile if the flag has been set. // Configure the workspace profile if the flag has been set.
diags = diags.Extend(configureProfile(cmd, b)) diags = diags.Extend(configureProfile(cmd, b))
if diags.HasError() {
return b, diags
}
// Set the auth configuration in the command context. This can be used
// downstream to initialize a API client.
//
// Note that just initializing a workspace client and loading auth configuration
// is a fast operation. It does not perform network I/O or invoke processes (for example the Azure CLI).
client, err := b.WorkspaceClientE()
if err != nil {
return b, diags.Extend(diag.FromErr(err))
}
ctx = context.WithValue(ctx, &configUsed, client.Config)
cmd.SetContext(ctx)
return b, diags return b, diags
} }

View File

@ -8,7 +8,6 @@ import (
"runtime" "runtime"
"testing" "testing"
"github.com/databricks/cli/bundle"
"github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/internal/testutil"
"github.com/spf13/cobra" "github.com/spf13/cobra"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
@ -38,7 +37,7 @@ func emptyCommand(t *testing.T) *cobra.Command {
return cmd return cmd
} }
func setupWithHost(t *testing.T, cmd *cobra.Command, host string) *bundle.Bundle { func setupWithHost(t *testing.T, cmd *cobra.Command, host string) error {
setupDatabricksCfg(t) setupDatabricksCfg(t)
rootPath := t.TempDir() rootPath := t.TempDir()
@ -51,12 +50,11 @@ workspace:
err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0o644) err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0o644)
require.NoError(t, err) require.NoError(t, err)
b, diags := MustConfigureBundle(cmd) _, diags := MustConfigureBundle(cmd)
require.NoError(t, diags.Error()) return diags.Error()
return b
} }
func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) *bundle.Bundle { func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) error {
setupDatabricksCfg(t) setupDatabricksCfg(t)
rootPath := t.TempDir() rootPath := t.TempDir()
@ -69,29 +67,25 @@ workspace:
err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0o644) err := os.WriteFile(filepath.Join(rootPath, "databricks.yml"), []byte(contents), 0o644)
require.NoError(t, err) require.NoError(t, err)
b, diags := MustConfigureBundle(cmd) _, diags := MustConfigureBundle(cmd)
require.NoError(t, diags.Error()) return diags.Error()
return b
} }
func TestBundleConfigureDefault(t *testing.T) { func TestBundleConfigureDefault(t *testing.T) {
testutil.CleanupEnvironment(t) testutil.CleanupEnvironment(t)
cmd := emptyCommand(t) cmd := emptyCommand(t)
b := setupWithHost(t, cmd, "https://x.com") err := setupWithHost(t, cmd, "https://x.com")
client, err := b.InitializeWorkspaceClient()
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://x.com", client.Config.Host)
assert.Equal(t, "https://x.com", ConfigUsed(cmd.Context()).Host)
} }
func TestBundleConfigureWithMultipleMatches(t *testing.T) { func TestBundleConfigureWithMultipleMatches(t *testing.T) {
testutil.CleanupEnvironment(t) testutil.CleanupEnvironment(t)
cmd := emptyCommand(t) cmd := emptyCommand(t)
b := setupWithHost(t, cmd, "https://a.com") err := setupWithHost(t, cmd, "https://a.com")
_, err := b.InitializeWorkspaceClient()
assert.ErrorContains(t, err, "multiple profiles matched: PROFILE-1, PROFILE-2") assert.ErrorContains(t, err, "multiple profiles matched: PROFILE-1, PROFILE-2")
} }
@ -101,9 +95,8 @@ func TestBundleConfigureWithNonExistentProfileFlag(t *testing.T) {
cmd := emptyCommand(t) cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("NOEXIST") err := cmd.Flag("profile").Value.Set("NOEXIST")
require.NoError(t, err) require.NoError(t, err)
b := setupWithHost(t, cmd, "https://x.com")
_, err = b.InitializeWorkspaceClient() err = setupWithHost(t, cmd, "https://x.com")
assert.ErrorContains(t, err, "has no NOEXIST profile configured") assert.ErrorContains(t, err, "has no NOEXIST profile configured")
} }
@ -113,9 +106,8 @@ func TestBundleConfigureWithMismatchedProfile(t *testing.T) {
cmd := emptyCommand(t) cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("PROFILE-1") err := cmd.Flag("profile").Value.Set("PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
b := setupWithHost(t, cmd, "https://x.com")
_, err = b.InitializeWorkspaceClient() err = setupWithHost(t, cmd, "https://x.com")
assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com") assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
} }
@ -125,12 +117,11 @@ func TestBundleConfigureWithCorrectProfile(t *testing.T) {
cmd := emptyCommand(t) cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("PROFILE-1") err := cmd.Flag("profile").Value.Set("PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
b := setupWithHost(t, cmd, "https://a.com") err = setupWithHost(t, cmd, "https://a.com")
client, err := b.InitializeWorkspaceClient()
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host) assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
assert.Equal(t, "PROFILE-1", client.Config.Profile) assert.Equal(t, "PROFILE-1", ConfigUsed(cmd.Context()).Profile)
} }
func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) { func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
@ -138,9 +129,8 @@ func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1") t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1")
cmd := emptyCommand(t) cmd := emptyCommand(t)
b := setupWithHost(t, cmd, "https://x.com")
_, err := b.InitializeWorkspaceClient() err := setupWithHost(t, cmd, "https://x.com")
assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com") assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
} }
@ -151,12 +141,11 @@ func TestBundleConfigureWithProfileFlagAndEnvVariable(t *testing.T) {
cmd := emptyCommand(t) cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("PROFILE-1") err := cmd.Flag("profile").Value.Set("PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
b := setupWithHost(t, cmd, "https://a.com")
client, err := b.InitializeWorkspaceClient() err = setupWithHost(t, cmd, "https://a.com")
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host) assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
assert.Equal(t, "PROFILE-1", client.Config.Profile) assert.Equal(t, "PROFILE-1", ConfigUsed(cmd.Context()).Profile)
} }
func TestBundleConfigureProfileDefault(t *testing.T) { func TestBundleConfigureProfileDefault(t *testing.T) {
@ -164,13 +153,12 @@ func TestBundleConfigureProfileDefault(t *testing.T) {
// The profile in the databricks.yml file is used // The profile in the databricks.yml file is used
cmd := emptyCommand(t) cmd := emptyCommand(t)
b := setupWithProfile(t, cmd, "PROFILE-1")
client, err := b.InitializeWorkspaceClient() err := setupWithProfile(t, cmd, "PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host) assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
assert.Equal(t, "a", client.Config.Token) assert.Equal(t, "a", ConfigUsed(cmd.Context()).Token)
assert.Equal(t, "PROFILE-1", client.Config.Profile) assert.Equal(t, "PROFILE-1", ConfigUsed(cmd.Context()).Profile)
} }
func TestBundleConfigureProfileFlag(t *testing.T) { func TestBundleConfigureProfileFlag(t *testing.T) {
@ -180,13 +168,12 @@ func TestBundleConfigureProfileFlag(t *testing.T) {
cmd := emptyCommand(t) cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("PROFILE-2") err := cmd.Flag("profile").Value.Set("PROFILE-2")
require.NoError(t, err) require.NoError(t, err)
b := setupWithProfile(t, cmd, "PROFILE-1")
client, err := b.InitializeWorkspaceClient() err = setupWithProfile(t, cmd, "PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host) assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
assert.Equal(t, "b", client.Config.Token) assert.Equal(t, "b", ConfigUsed(cmd.Context()).Token)
assert.Equal(t, "PROFILE-2", client.Config.Profile) assert.Equal(t, "PROFILE-2", ConfigUsed(cmd.Context()).Profile)
} }
func TestBundleConfigureProfileEnvVariable(t *testing.T) { func TestBundleConfigureProfileEnvVariable(t *testing.T) {
@ -195,13 +182,12 @@ func TestBundleConfigureProfileEnvVariable(t *testing.T) {
// The DATABRICKS_CONFIG_PROFILE environment variable takes precedence over the profile in the databricks.yml file // The DATABRICKS_CONFIG_PROFILE environment variable takes precedence over the profile in the databricks.yml file
t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-2") t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-2")
cmd := emptyCommand(t) cmd := emptyCommand(t)
b := setupWithProfile(t, cmd, "PROFILE-1")
client, err := b.InitializeWorkspaceClient() err := setupWithProfile(t, cmd, "PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host) assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
assert.Equal(t, "b", client.Config.Token) assert.Equal(t, "b", ConfigUsed(cmd.Context()).Token)
assert.Equal(t, "PROFILE-2", client.Config.Profile) assert.Equal(t, "PROFILE-2", ConfigUsed(cmd.Context()).Profile)
} }
func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) { func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) {
@ -212,13 +198,12 @@ func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) {
cmd := emptyCommand(t) cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("PROFILE-2") err := cmd.Flag("profile").Value.Set("PROFILE-2")
require.NoError(t, err) require.NoError(t, err)
b := setupWithProfile(t, cmd, "PROFILE-1")
client, err := b.InitializeWorkspaceClient() err = setupWithProfile(t, cmd, "PROFILE-1")
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host) assert.Equal(t, "https://a.com", ConfigUsed(cmd.Context()).Host)
assert.Equal(t, "b", client.Config.Token) assert.Equal(t, "b", ConfigUsed(cmd.Context()).Token)
assert.Equal(t, "PROFILE-2", client.Config.Profile) assert.Equal(t, "PROFILE-2", ConfigUsed(cmd.Context()).Profile)
} }
func TestTargetFlagFull(t *testing.T) { func TestTargetFlagFull(t *testing.T) {