Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Set bundle auth configuration in command context #2195

Merged
merged 16 commits into from
Jan 29, 2025
Merged

Set bundle auth configuration in command context #2195

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
bf68318
Set bundle auth configuration in command context
shreyas-goenka Jan 20, 2025
e868dbc
fix build
shreyas-goenka Jan 20, 2025
27f1ae3
fix tests
shreyas-goenka Jan 20, 2025
9cff7c6
Revert "fix tests"
shreyas-goenka Jan 20, 2025
506a4bd
fix tests easy to review version
shreyas-goenka Jan 20, 2025
bebbf4b
Merge remote-tracking branch 'origin' into set-config-used-bundle
shreyas-goenka Jan 21, 2025
8cc50c2
set client during initialization
shreyas-goenka Jan 21, 2025
c0cf3f8
remove unused parameter
shreyas-goenka Jan 21, 2025
79a52dc
Revert "remove unused parameter"
shreyas-goenka Jan 21, 2025
c07f137
Revert "set client during initialization"
shreyas-goenka Jan 21, 2025
7855d91
modify approach
shreyas-goenka Jan 21, 2025
aae8e18
Merge remote-tracking branch 'origin' into set-config-used-bundle
shreyas-goenka Jan 22, 2025
d147308
fix infinite recursion
shreyas-goenka Jan 22, 2025
dc5909f
Merge remote-tracking branch 'origin' into set-config-used-bundle
shreyas-goenka Jan 28, 2025
334ceda
add comment
shreyas-goenka Jan 28, 2025
7f33e95
Update cmd/root/bundle.go
shreyas-goenka Jan 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 0 additions & 26 deletions bundle/config/mutator/initialize_workspace_client.go
View file
Open in desktop

This file was deleted.

1 change: 0 additions & 1 deletion bundle/phases/initialize.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),

Expand Down
13 changes: 13 additions & 0 deletions cmd/root/bundle.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,19 @@ 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.
client, err := b.InitializeWorkspaceClient()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is also WorkspaceClient() which creates and caches the client on bundle.

If that's not necessary, should we clean that up?

Related: Bundle.client, Bundke.clientOnce, can we get rid of those?

Copy link
Contributor Author

@shreyas-goenka shreyas-goenka Jan 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If that's not necessary, should we clean that up?

Eyeballing it we could clean that up a bit without issues. It's out of the scope of what this PR is trying to achieve, though.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can send a follow-up PR to do so.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's in scope, we should not have multiple ways to set up client.

Also, please take a look "variables v2" doc. How would this work if we want to interpolate variables first?

Copy link
Contributor Author

@shreyas-goenka shreyas-goenka Jan 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should not have multiple ways to set up client.

Fair enough. I modified the change to set the client as well. I was hesitant to do this because I don't like how decoupled the client's setting is from its usage, but this change should be safe enough (and covered by tests).

How would this work if we want to interpolate variables first?

We'd need to resolve the variables before we resolve the client.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed the approach to still ensure that the client is only initialized / set once. Did so by renaming InitializeWorkspaceClient to WorkspaceClientE

if err != nil {
return b, diags.Extend(diag.FromErr(err))
}
ctx = context.WithValue(ctx, &configUsed, client.Config)
cmd.SetContext(ctx)

return b, diags
}

Expand Down
312 changes: 157 additions & 155 deletions cmd/root/bundle_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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) {
setupDatabricksCfg(t)

rootPath := t.TempDir()
Expand All @@ -50,13 +49,9 @@ workspace:
`, host)
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
}

func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) *bundle.Bundle {
func setupWithProfile(t *testing.T, cmd *cobra.Command, profile string) {
setupDatabricksCfg(t)

rootPath := t.TempDir()
Expand All @@ -68,157 +63,164 @@ workspace:
`, profile)
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
}

func TestBundleConfigureDefault(t *testing.T) {
testutil.CleanupEnvironment(t)

cmd := emptyCommand(t)
b := setupWithHost(t, cmd, "https://x.com")

client, err := b.InitializeWorkspaceClient()
require.NoError(t, err)
assert.Equal(t, "https://x.com", client.Config.Host)
}

func TestBundleConfigureWithMultipleMatches(t *testing.T) {
testutil.CleanupEnvironment(t)

cmd := emptyCommand(t)
b := setupWithHost(t, cmd, "https://a.com")

_, err := b.InitializeWorkspaceClient()
assert.ErrorContains(t, err, "multiple profiles matched: PROFILE-1, PROFILE-2")
}

func TestBundleConfigureWithNonExistentProfileFlag(t *testing.T) {
testutil.CleanupEnvironment(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()
assert.ErrorContains(t, err, "has no NOEXIST profile configured")
}

func TestBundleConfigureWithMismatchedProfile(t *testing.T) {
testutil.CleanupEnvironment(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()
assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
}

func TestBundleConfigureWithCorrectProfile(t *testing.T) {
testutil.CleanupEnvironment(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()
require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host)
assert.Equal(t, "PROFILE-1", client.Config.Profile)
}

func TestBundleConfigureWithMismatchedProfileEnvVariable(t *testing.T) {
testutil.CleanupEnvironment(t)

t.Setenv("DATABRICKS_CONFIG_PROFILE", "PROFILE-1")
cmd := emptyCommand(t)
b := setupWithHost(t, cmd, "https://x.com")

_, err := b.InitializeWorkspaceClient()
assert.ErrorContains(t, err, "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com")
}

func TestBundleConfigureWithProfileFlagAndEnvVariable(t *testing.T) {
testutil.CleanupEnvironment(t)

t.Setenv("DATABRICKS_CONFIG_PROFILE", "NOEXIST")
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()
require.NoError(t, err)
assert.Equal(t, "https://a.com", client.Config.Host)
assert.Equal(t, "PROFILE-1", client.Config.Profile)
}

func TestBundleConfigureProfileDefault(t *testing.T) {
testutil.CleanupEnvironment(t)

// The profile in the databricks.yml file is used
cmd := emptyCommand(t)
b := setupWithProfile(t, cmd, "PROFILE-1")

client, err := b.InitializeWorkspaceClient()
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)
}

func TestBundleConfigureProfileFlag(t *testing.T) {
testutil.CleanupEnvironment(t)

// The --profile flag takes precedence over the profile in the databricks.yml file
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()
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)
}

func TestBundleConfigureProfileEnvVariable(t *testing.T) {
testutil.CleanupEnvironment(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()
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)
}

func TestBundleConfigureProfileFlagAndEnvVariable(t *testing.T) {
testutil.CleanupEnvironment(t)

// The --profile flag takes precedence over the DATABRICKS_CONFIG_PROFILE environment variable
t.Setenv("DATABRICKS_CONFIG_PROFILE", "NOEXIST")
cmd := emptyCommand(t)
err := cmd.Flag("profile").Value.Set("PROFILE-2")
require.NoError(t, err)
b := setupWithProfile(t, cmd, "PROFILE-1")
func TestBundleConfigureProfile(t *testing.T) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While acceptance tests would be better suited for most of these, since this PR changes how auth is loaded I opted to keep the unit tests.

tcases := []struct {
name string
hostInConfig string

// --profile flag
profileFlag string
// DATABRICKS_CONFIG_PROFILE environment variable
profileEnvVar string
// profile in config
profileInConfig string

expectedError string
expectedHost string
expectedProfile string
expectedToken string
}{
{
name: "no match, keep host",
hostInConfig: "https://x.com",

expectedHost: "https://x.com",
},
{
name: "multiple profile matches",
hostInConfig: "https://a.com",

expectedError: "multiple profiles matched: PROFILE-1, PROFILE-2",
},
{
name: "non-existent profile",
profileFlag: "NOEXIST",
hostInConfig: "https://x.com",

expectedError: "has no NOEXIST profile configured",
},
{
name: "mismatched profile",
hostInConfig: "https://x.com",
profileFlag: "PROFILE-1",

expectedError: "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com",
},
{
name: "profile flag specified",
hostInConfig: "https://a.com",
profileFlag: "PROFILE-1",

expectedHost: "https://a.com",
expectedProfile: "PROFILE-1",
},
{
name: "mismatched profile env variable",
hostInConfig: "https://x.com",
profileEnvVar: "PROFILE-1",

expectedError: "config host mismatch: profile uses host https://a.com, but CLI configured to use https://x.com",
},
{
// The --profile flag takes precedence over the DATABRICKS_CONFIG_PROFILE environment variable
name: "(host) profile flag takes precedence over env variable",
hostInConfig: "https://a.com",
profileFlag: "PROFILE-1",
profileEnvVar: "NOEXIST",

expectedHost: "https://a.com",
expectedProfile: "PROFILE-1",
},
{
name: "profile from config",
profileInConfig: "PROFILE-1",

expectedHost: "https://a.com",
expectedProfile: "PROFILE-1",
expectedToken: "a",
},
{
// The --profile flag takes precedence over the profile in the databricks.yml file
name: "profile flag takes precedence",
profileInConfig: "PROFILE-1",
profileFlag: "PROFILE-2",

expectedHost: "https://a.com",
expectedProfile: "PROFILE-2",
expectedToken: "b",
},
{
// The DATABRICKS_CONFIG_PROFILE environment variable takes precedence over the profile in the databricks.yml file
name: "profile env variable takes precedence",
profileInConfig: "PROFILE-1",
profileEnvVar: "PROFILE-2",

expectedHost: "https://a.com",
expectedProfile: "PROFILE-2",
expectedToken: "b",
},
{
// The --profile flag takes precedence over the DATABRICKS_CONFIG_PROFILE environment variable
name: "profile flag takes precedence over env variable",
profileInConfig: "PROFILE-1",
profileFlag: "PROFILE-2",
profileEnvVar: "NOEXIST",

expectedHost: "https://a.com",
expectedProfile: "PROFILE-2",
expectedToken: "b",
},
}

client, err := b.InitializeWorkspaceClient()
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)
for _, tc := range tcases {
t.Run(tc.name, func(t *testing.T) {
testutil.CleanupEnvironment(t)

cmd := emptyCommand(t)

// Set up host in databricks.yml
if tc.hostInConfig != "" {
setupWithHost(t, cmd, tc.hostInConfig)
}

// Set up profile in databricks.yml
if tc.profileInConfig != "" {
setupWithProfile(t, cmd, tc.profileInConfig)
}

// Set --profile flag
if tc.profileFlag != "" {
err := cmd.Flag("profile").Value.Set(tc.profileFlag)
require.NoError(t, err)
}

// Set DATABRICKS_CONFIG_PROFILE environment variable
if tc.profileEnvVar != "" {
t.Setenv("DATABRICKS_CONFIG_PROFILE", tc.profileEnvVar)
}

_, diags := MustConfigureBundle(cmd)

if tc.expectedError != "" {
assert.ErrorContains(t, diags.Error(), tc.expectedError)
} else {
assert.NoError(t, diags.Error())
}

// Assert on the resolved configuration values
if tc.expectedHost != "" {
assert.Equal(t, tc.expectedHost, ConfigUsed(cmd.Context()).Host)
}
if tc.expectedProfile != "" {
assert.Equal(t, tc.expectedProfile, ConfigUsed(cmd.Context()).Profile)
}
if tc.expectedToken != "" {
assert.Equal(t, tc.expectedToken, ConfigUsed(cmd.Context()).Token)
}
})
}
}

func TestTargetFlagFull(t *testing.T) {
Expand Down
Loading