From bb6eadc23899fd982f4e16ede6835c159f9aad8a Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 10 Oct 2023 13:55:09 +0200 Subject: [PATCH 1/3] Use profile information when getting a token using the CLI --- cmd/auth/token.go | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 242a3dabee..12d305e76c 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -6,6 +6,7 @@ import ( "time" "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/databrickscfg" "github.com/spf13/cobra" ) @@ -21,9 +22,27 @@ func newTokenCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() + + var profileName string + profileFlag := cmd.Flag("profile") + if profileFlag != nil { + profileName = profileFlag.Value.String() + } + // If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile. + _, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool { + return p.Name == profileName + }) + if err != nil { + return err + } if persistentAuth.Host == "" { - configureHost(ctx, persistentAuth, args, 0) + if len(profiles) > 0 && profiles[0].Host != "" { + persistentAuth.Host = profiles[0].Host + } else { + configureHost(ctx, persistentAuth, args, 0) + } } + defer persistentAuth.Close() ctx, cancel := context.WithTimeout(ctx, tokenTimeout) From 3a38e536ec912f6dd7bb536dec351f65166f6cd9 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Tue, 10 Oct 2023 15:28:53 +0200 Subject: [PATCH 2/3] Extract method to set host --- cmd/auth/login.go | 30 +++++++++++++++++++----------- cmd/auth/token.go | 15 ++------------- 2 files changed, 21 insertions(+), 24 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index a14c5ebefc..3a3f3a6dcf 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -60,20 +60,10 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { profileName = profile } - // If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile. - _, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool { - return p.Name == profileName - }) + err := setHost(ctx, profileName, persistentAuth, args) if err != nil { return err } - if persistentAuth.Host == "" { - if len(profiles) > 0 && profiles[0].Host != "" { - persistentAuth.Host = profiles[0].Host - } else { - configureHost(ctx, persistentAuth, args, 0) - } - } defer persistentAuth.Close() // We need the config without the profile before it's used to initialise new workspace client below. @@ -135,3 +125,21 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { return cmd } + +func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error { + // If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile. + _, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool { + return p.Name == profileName + }) + if err != nil { + return err + } + if persistentAuth.Host == "" { + if len(profiles) > 0 && profiles[0].Host != "" { + persistentAuth.Host = profiles[0].Host + } else { + configureHost(ctx, persistentAuth, args, 0) + } + } + return nil +} diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 12d305e76c..3bc39989b0 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -6,7 +6,6 @@ import ( "time" "github.com/databricks/cli/libs/auth" - "github.com/databricks/cli/libs/databrickscfg" "github.com/spf13/cobra" ) @@ -28,21 +27,11 @@ func newTokenCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { if profileFlag != nil { profileName = profileFlag.Value.String() } - // If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile. - _, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool { - return p.Name == profileName - }) + + err := setHost(ctx, profileName, persistentAuth, args) if err != nil { return err } - if persistentAuth.Host == "" { - if len(profiles) > 0 && profiles[0].Host != "" { - persistentAuth.Host = profiles[0].Host - } else { - configureHost(ctx, persistentAuth, args, 0) - } - } - defer persistentAuth.Close() ctx, cancel := context.WithTimeout(ctx, tokenTimeout) From d43262789d2f4f576068ddf849b0526f37705049 Mon Sep 17 00:00:00 2001 From: Hector Castejon Diaz Date: Wed, 11 Oct 2023 09:14:54 +0200 Subject: [PATCH 3/3] Throw an error if both profile and host are provided --- cmd/auth/token.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 3bc39989b0..d763b95642 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -3,6 +3,7 @@ package auth import ( "context" "encoding/json" + "errors" "time" "github.com/databricks/cli/libs/auth" @@ -26,6 +27,10 @@ func newTokenCommand(persistentAuth *auth.PersistentAuth) *cobra.Command { profileFlag := cmd.Flag("profile") if profileFlag != nil { profileName = profileFlag.Value.String() + // If a profile is provided we read the host from the .databrickscfg file + if profileName != "" && len(args) > 0 { + return errors.New("providing both a profile and a host parameters is not supported") + } } err := setHost(ctx, profileName, persistentAuth, args)