From 5ed635a24091872d00d2a20ee9cefcf829a556ac Mon Sep 17 00:00:00 2001
From: Serge Smertin <259697+nfx@users.noreply.github.com>
Date: Mon, 21 Aug 2023 18:17:02 +0200
Subject: [PATCH] Added `databricks account o-auth-enrollment enable` command
 (#687)

This command takes the user through the interactive flow to set up OAuth
for a fresh account, where only Basic authentication works.

---------

Co-authored-by: Andrew Nester <andrew.nester@databricks.com>
---
 cmd/account/o-auth-enrollment/overrides.go | 107 +++++++++++++++++++++
 libs/cmdio/io.go                           |  29 ++++++
 2 files changed, 136 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 00000000..1fc3aacc
--- /dev/null
+++ b/cmd/account/o-auth-enrollment/overrides.go
@@ -0,0 +1,107 @@
+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("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("failed to create oauth enrollment: %w", err)
+	}
+	enableSpinner := cmdio.Spinner(ctx)
+	// The actual enrollment take a few minutes
+	err = 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 := "Enabling OAuth..."
+			enableSpinner <- msg
+			return retries.Continues(msg)
+		}
+		enableSpinner <- "OAuth is enabled"
+		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",
+	})
+	if err != nil {
+		return fmt.Errorf("failed to enable databricks CLI: %w", err)
+	}
+	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 bc5a5f30..9d712e35 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 {