From 6875bec0e48496650c0e14cb6a33692d5f361e83 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Mon, 21 Aug 2023 16:55:48 +0200 Subject: [PATCH 1/4] Added `databricks account o-auth-enrollment enable` command This command takes user through interactive flow to setup OAuth for a fresh account, where only Basic authentication works. --- cmd/account/o-auth-enrollment/overrides.go | 93 ++++++++++++++++++++++ libs/cmdio/io.go | 29 +++++++ 2 files changed, 122 insertions(+) create mode 100644 cmd/account/o-auth-enrollment/overrides.go diff --git a/cmd/account/o-auth-enrollment/overrides.go b/cmd/account/o-auth-enrollment/overrides.go new file mode 100644 index 0000000000..d37fd470b2 --- /dev/null +++ b/cmd/account/o-auth-enrollment/overrides.go @@ -0,0 +1,93 @@ +package o_auth_enrollment + +import ( + "context" + "fmt" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/retries" + "github.com/databricks/databricks-sdk-go/service/oauth2" + "github.com/spf13/cobra" +) + +func promptForBasicAccountConfig(ctx context.Context) (*databricks.Config, error) { + if !cmdio.IsInTTY(ctx) { + return nil, fmt.Errorf("this command requires a TTY") + } + // OAuth Enrollment only works on AWS + host, err := cmdio.DefaultPrompt(ctx, "Host", "https://accounts.cloud.databricks.com") + if err != nil { + return nil, fmt.Errorf("host: %w", err) + } + accountID, err := cmdio.SimplePrompt(ctx, "Account ID") + if err != nil { + return nil, fmt.Errorf("account: %w", err) + } + username, err := cmdio.SimplePrompt(ctx, "Username") + if err != nil { + return nil, fmt.Errorf("username: %w", err) + } + password, err := cmdio.Secret(ctx, "Password") + if err != nil { + return nil, fmt.Errorf("password: %w", err) + } + return &databricks.Config{ + Host: host, + AccountID: accountID, + Username: username, + Password: password, + }, nil +} + +func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { + ac, err := databricks.NewAccountClient(cfg) + if err != nil { + return fmt.Errorf("account client: %w", err) + } + // The enrollment is executed asynchronously, so the API returns HTTP 204 immediately + err = ac.OAuthEnrollment.Create(ctx, oauth2.CreateOAuthEnrollment{ + EnableAllPublishedApps: true, + }) + enableSpinner := cmdio.Spinner(ctx) + // The actual enrollment take a few minutes + return retries.Wait(ctx, 10*time.Minute, func() *retries.Err { + status, err := ac.OAuthEnrollment.Get(ctx) + if err != nil { + return retries.Halt(err) + } + if !status.IsEnabled { + msg := "OAuth is not yet enalbed" + enableSpinner <- msg + return retries.Continues(msg) + } + enableSpinner <- "OAuth is enabled" + close(enableSpinner) + return nil + }) +} + +func newEnable() *cobra.Command { + return &cobra.Command{ + Use: "enable", + Short: "Enable Databricks CLI, Tableau Desktop, and PowerBI for this account.", + Long: `Before you can do 'databricks auth login', you have to enable OAuth for this account. + +This command prompts you for Account ID, username, and password and waits until OAuth is enabled.`, + RunE: func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + cfg, err := promptForBasicAccountConfig(ctx) + if err != nil { + return fmt.Errorf("account config: %w", err) + } + return enableOAuthForAccount(ctx, cfg) + }, + } +} + +func init() { + cmdOverrides = append(cmdOverrides, func(c *cobra.Command) { + c.AddCommand(newEnable()) + }) +} diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index bc5a5f302b..9d712e351a 100644 --- a/libs/cmdio/io.go +++ b/libs/cmdio/io.go @@ -205,6 +205,35 @@ func Prompt(ctx context.Context) *promptui.Prompt { } } +func (c *cmdIO) simplePrompt(label string) *promptui.Prompt { + return &promptui.Prompt{ + Label: label, + Stdin: io.NopCloser(c.in), + Stdout: nopWriteCloser{c.out}, + } +} + +func (c *cmdIO) SimplePrompt(label string) (value string, err error) { + return c.simplePrompt(label).Run() +} + +func SimplePrompt(ctx context.Context, label string) (value string, err error) { + c := fromContext(ctx) + return c.SimplePrompt(label) +} + +func (c *cmdIO) DefaultPrompt(label, defaultValue string) (value string, err error) { + prompt := c.simplePrompt(label) + prompt.Default = defaultValue + prompt.AllowEdit = true + return prompt.Run() +} + +func DefaultPrompt(ctx context.Context, label, defaultValue string) (value string, err error) { + c := fromContext(ctx) + return c.DefaultPrompt(label, defaultValue) +} + func (c *cmdIO) Spinner(ctx context.Context) chan string { var sp *spinner.Spinner if c.interactive { From bfe9808b255db162aa2fc9efff1ed79669a7ec33 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Mon, 21 Aug 2023 17:10:33 +0200 Subject: [PATCH 2/4] add published app --- cmd/account/o-auth-enrollment/overrides.go | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/cmd/account/o-auth-enrollment/overrides.go b/cmd/account/o-auth-enrollment/overrides.go index d37fd470b2..d763a62331 100644 --- a/cmd/account/o-auth-enrollment/overrides.go +++ b/cmd/account/o-auth-enrollment/overrides.go @@ -50,9 +50,12 @@ func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { err = ac.OAuthEnrollment.Create(ctx, oauth2.CreateOAuthEnrollment{ EnableAllPublishedApps: true, }) + if err != nil { + return fmt.Errorf("enroll: %w", err) + } enableSpinner := cmdio.Spinner(ctx) // The actual enrollment take a few minutes - return retries.Wait(ctx, 10*time.Minute, func() *retries.Err { + err = retries.Wait(ctx, 10*time.Minute, func() *retries.Err { status, err := ac.OAuthEnrollment.Get(ctx) if err != nil { return retries.Halt(err) @@ -66,6 +69,14 @@ func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { close(enableSpinner) return nil }) + // enable Databricks CLI, so that `databricks auth login` works + _, err = ac.PublishedAppIntegration.Create(ctx, oauth2.CreatePublishedAppIntegration{ + AppId: "databricks-cli", + }) + if err != nil { + return fmt.Errorf("enabling databricks-cli: %w", err) + } + return nil } func newEnable() *cobra.Command { From 574ab34e77da6315175ec65f232b576325ba56ba Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Mon, 21 Aug 2023 17:19:08 +0200 Subject: [PATCH 3/4] fix lint --- cmd/account/o-auth-enrollment/overrides.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cmd/account/o-auth-enrollment/overrides.go b/cmd/account/o-auth-enrollment/overrides.go index d763a62331..272ad82d69 100644 --- a/cmd/account/o-auth-enrollment/overrides.go +++ b/cmd/account/o-auth-enrollment/overrides.go @@ -69,6 +69,9 @@ func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { close(enableSpinner) return nil }) + if err != nil { + return fmt.Errorf("wait for enrollment: %w", err) + } // enable Databricks CLI, so that `databricks auth login` works _, err = ac.PublishedAppIntegration.Create(ctx, oauth2.CreatePublishedAppIntegration{ AppId: "databricks-cli", From fc7f0220c8a6450d03a9956efeec018d53d96455 Mon Sep 17 00:00:00 2001 From: Serge Smertin <259697+nfx@users.noreply.github.com> Date: Mon, 21 Aug 2023 17:21:55 +0200 Subject: [PATCH 4/4] Apply suggestions from code review Co-authored-by: Andrew Nester --- cmd/account/o-auth-enrollment/overrides.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cmd/account/o-auth-enrollment/overrides.go b/cmd/account/o-auth-enrollment/overrides.go index 272ad82d69..1fc3aacc1d 100644 --- a/cmd/account/o-auth-enrollment/overrides.go +++ b/cmd/account/o-auth-enrollment/overrides.go @@ -44,14 +44,14 @@ func promptForBasicAccountConfig(ctx context.Context) (*databricks.Config, error func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { ac, err := databricks.NewAccountClient(cfg) if err != nil { - return fmt.Errorf("account client: %w", err) + return fmt.Errorf("failed to instantiate account client: %w", err) } // The enrollment is executed asynchronously, so the API returns HTTP 204 immediately err = ac.OAuthEnrollment.Create(ctx, oauth2.CreateOAuthEnrollment{ EnableAllPublishedApps: true, }) if err != nil { - return fmt.Errorf("enroll: %w", err) + return fmt.Errorf("failed to create oauth enrollment: %w", err) } enableSpinner := cmdio.Spinner(ctx) // The actual enrollment take a few minutes @@ -61,7 +61,7 @@ func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { return retries.Halt(err) } if !status.IsEnabled { - msg := "OAuth is not yet enalbed" + msg := "Enabling OAuth..." enableSpinner <- msg return retries.Continues(msg) } @@ -77,7 +77,7 @@ func enableOAuthForAccount(ctx context.Context, cfg *databricks.Config) error { AppId: "databricks-cli", }) if err != nil { - return fmt.Errorf("enabling databricks-cli: %w", err) + return fmt.Errorf("failed to enable databricks CLI: %w", err) } return nil }