From 124515e8d2105a3d2ec071dadcb30bf792ba9cad Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Wed, 29 Jan 2025 16:12:21 +0530 Subject: [PATCH 1/2] Move TestServer from acceptance to libs/testserver (#2255) ## Changes Just a move, no changes. As recommended here: https://github.com/databricks/cli/pull/2226#discussion_r1932152627 ## Tests N/A --- acceptance/acceptance_test.go | 3 +- acceptance/cmd_server_test.go | 3 +- acceptance/server_test.go | 56 +++---------------------------- libs/testserver/server.go | 63 +++++++++++++++++++++++++++++++++++ 4 files changed, 71 insertions(+), 54 deletions(-) create mode 100644 libs/testserver/server.go diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index b4b27f201..91ad09e9e 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -20,6 +20,7 @@ import ( "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testdiff" + "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" "github.com/stretchr/testify/require" ) @@ -107,7 +108,7 @@ func testAccept(t *testing.T, InprocessMode bool, singleTest string) int { cloudEnv := os.Getenv("CLOUD_ENV") if cloudEnv == "" { - server := StartServer(t) + server := testserver.New(t) AddHandlers(server) // Redirect API access to local server: t.Setenv("DATABRICKS_HOST", server.URL) diff --git a/acceptance/cmd_server_test.go b/acceptance/cmd_server_test.go index 28feec1bd..3f5a6356e 100644 --- a/acceptance/cmd_server_test.go +++ b/acceptance/cmd_server_test.go @@ -8,10 +8,11 @@ import ( "testing" "github.com/databricks/cli/internal/testcli" + "github.com/databricks/cli/libs/testserver" "github.com/stretchr/testify/require" ) -func StartCmdServer(t *testing.T) *TestServer { +func StartCmdServer(t *testing.T) *testserver.Server { server := StartServer(t) server.Handle("/", func(r *http.Request) (any, error) { q := r.URL.Query() diff --git a/acceptance/server_test.go b/acceptance/server_test.go index dbc55c03f..66de5dcbf 100644 --- a/acceptance/server_test.go +++ b/acceptance/server_test.go @@ -1,73 +1,25 @@ package acceptance_test import ( - "encoding/json" "net/http" - "net/http/httptest" "testing" + "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/workspace" ) -type TestServer struct { - *httptest.Server - Mux *http.ServeMux -} - -type HandlerFunc func(r *http.Request) (any, error) - -func NewTestServer() *TestServer { - mux := http.NewServeMux() - server := httptest.NewServer(mux) - - return &TestServer{ - Server: server, - Mux: mux, - } -} - -func (s *TestServer) Handle(pattern string, handler HandlerFunc) { - s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) { - resp, err := handler(r) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - - w.Header().Set("Content-Type", "application/json") - - var respBytes []byte - - respString, ok := resp.(string) - if ok { - respBytes = []byte(respString) - } else { - respBytes, err = json.MarshalIndent(resp, "", " ") - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - } - - if _, err := w.Write(respBytes); err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - }) -} - -func StartServer(t *testing.T) *TestServer { - server := NewTestServer() +func StartServer(t *testing.T) *testserver.Server { + server := testserver.New(t) t.Cleanup(func() { server.Close() }) return server } -func AddHandlers(server *TestServer) { +func AddHandlers(server *testserver.Server) { server.Handle("GET /api/2.0/policies/clusters/list", func(r *http.Request) (any, error) { return compute.ListPoliciesResponse{ Policies: []compute.Policy{ diff --git a/libs/testserver/server.go b/libs/testserver/server.go new file mode 100644 index 000000000..10269af8f --- /dev/null +++ b/libs/testserver/server.go @@ -0,0 +1,63 @@ +package testserver + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + + "github.com/databricks/cli/internal/testutil" +) + +type Server struct { + *httptest.Server + Mux *http.ServeMux + + t testutil.TestingT +} + +func New(t testutil.TestingT) *Server { + mux := http.NewServeMux() + server := httptest.NewServer(mux) + + return &Server{ + Server: server, + Mux: mux, + t: t, + } +} + +type HandlerFunc func(req *http.Request) (resp any, err error) + +func (s *Server) Close() { + s.Server.Close() +} + +func (s *Server) Handle(pattern string, handler HandlerFunc) { + s.Mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) { + resp, err := handler(r) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "application/json") + + var respBytes []byte + + respString, ok := resp.(string) + if ok { + respBytes = []byte(respString) + } else { + respBytes, err = json.MarshalIndent(resp, "", " ") + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + } + + if _, err := w.Write(respBytes); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + }) +} From 884b5f26ed148c431a0dfea6333bff9b293f8ed1 Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Wed, 29 Jan 2025 16:32:08 +0530 Subject: [PATCH 2/2] 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: https://github.com/databricks/cli/blob/41bbd89257285707b3c3df9b9e5b92d6bcf8f1d1/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`: https://github.com/databricks/cli/blob/41bbd89257285707b3c3df9b9e5b92d6bcf8f1d1/cmd/root/bundle.go#L88 2. `utils.ConfigureBundleWithVariables`: https://github.com/databricks/cli/blob/41bbd89257285707b3c3df9b9e5b92d6bcf8f1d1/cmd/bundle/utils/utils.go#L19 --------- Co-authored-by: Pieter Noordhuis --- bundle/bundle.go | 31 ++++--- .../mutator/initialize_workspace_client.go | 26 ------ bundle/phases/initialize.go | 1 - cmd/root/auth.go | 2 +- cmd/root/bundle.go | 16 ++++ cmd/root/bundle_test.go | 85 ++++++++----------- 6 files changed, 69 insertions(+), 92 deletions(-) delete mode 100644 bundle/config/mutator/initialize_workspace_client.go diff --git a/bundle/bundle.go b/bundle/bundle.go index e715b8b2c..9cb8916f5 100644 --- a/bundle/bundle.go +++ b/bundle/bundle.go @@ -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() - if err != nil { - return nil, fmt.Errorf("cannot resolve bundle auth configuration: %w", err) - } - return client, nil +func (b *Bundle) WorkspaceClientE() (*databricks.WorkspaceClient, error) { + b.clientOnce.Do(func() { + var err error + b.client, err = b.Config.Workspace.Client() + if err != nil { + b.clientErr = fmt.Errorf("cannot resolve bundle auth configuration: %w", err) + } + }) + + return b.client, b.clientErr } func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient { - b.clientOnce.Do(func() { - var err error - b.client, err = b.InitializeWorkspaceClient() - if err != nil { - panic(err) - } - }) - return b.client + client, err := b.WorkspaceClientE() + if err != nil { + panic(err) + } + + return client } // SetWorkpaceClient sets the workspace client for this bundle. diff --git a/bundle/config/mutator/initialize_workspace_client.go b/bundle/config/mutator/initialize_workspace_client.go deleted file mode 100644 index 5c905f40c..000000000 --- a/bundle/config/mutator/initialize_workspace_client.go +++ /dev/null @@ -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) -} diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index c5b875196..afd6def3f 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -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(), diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 49abfd414..4fcfbb4d8 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -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 } diff --git a/cmd/root/bundle.go b/cmd/root/bundle.go index 8b98f2cf2..5842526f3 100644 --- a/cmd/root/bundle.go +++ b/cmd/root/bundle.go @@ -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 } diff --git a/cmd/root/bundle_test.go b/cmd/root/bundle_test.go index 1998b19e6..3517b02e4 100644 --- a/cmd/root/bundle_test.go +++ b/cmd/root/bundle_test.go @@ -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) {