steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Set bundle auth configuration in command context (#2195) 

## Changes

This change is required to enable tracking execution time telemetry for
bundle commands. In order to track execution time for the command
generally, we need to have the databricks auth configuration available
at this section of the code:


41bbd89257/cmd/root/root.go (L99)

In order to do this we can rely on the `configUsed` context key.   

Most commands rely on the `root.MustWorkspaceClient` function which
automatically sets the client config in the `configUsed` context key.
Bundle commands, however, do not do so. They instead store their
workspace clients in the `&bundle.Bundle{}` object.

With this PR, the `configUsed` context key will be set for all `bundle`
commands. Functionally nothing changes.

## Tests
Existing tests. Also manually verified that either
`root.MustConfigureBundle` or `utils.ConfigureBundleWithVariables` is
called for all bundle commands (except `bundle init`) thus ensuring this
context key would be set for all bundle commands.

refs for the functions:
1. `root.MustConfigureBundle`:
41bbd89257/cmd/root/bundle.go (L88)
2. `utils.ConfigureBundleWithVariables`:
41bbd89257/cmd/bundle/utils/utils.go (L19)

---------

Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>
This commit is contained in:
shreyas-goenka 2025-01-29 16:32:08 +05:30 committed by GitHub
parent 124515e8d2
commit 884b5f26ed
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 69 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
bundle/bundle.go
View File

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

26
bundle/config/mutator/initialize_workspace_client.go
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)
}

1
bundle/phases/initialize.go
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.
mutator.SyncInferRoot(),
mutator.InitializeWorkspaceClient(),
mutator.PopulateCurrentUser(),
mutator.LoadGitDetails(),

2
cmd/root/auth.go
View File

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

16
cmd/root/bundle.go
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.
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
}

85
cmd/root/bundle_test.go
View File

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