2023-01-06 15:15:57 +00:00
|
|
|
package auth
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2023-11-23 09:04:54 +00:00
|
|
|
"errors"
|
2023-06-21 10:58:28 +00:00
|
|
|
"fmt"
|
2023-01-06 15:15:57 +00:00
|
|
|
"time"
|
|
|
|
|
2023-05-16 16:35:39 +00:00
|
|
|
"github.com/databricks/cli/libs/auth"
|
2023-06-21 10:58:28 +00:00
|
|
|
"github.com/databricks/cli/libs/cmdio"
|
|
|
|
"github.com/databricks/cli/libs/databrickscfg"
|
2023-11-09 16:38:45 +00:00
|
|
|
"github.com/databricks/cli/libs/databrickscfg/cfgpickers"
|
2023-06-21 15:51:59 +00:00
|
|
|
"github.com/databricks/databricks-sdk-go"
|
2023-06-21 10:58:28 +00:00
|
|
|
"github.com/databricks/databricks-sdk-go/config"
|
2023-01-06 15:15:57 +00:00
|
|
|
"github.com/spf13/cobra"
|
|
|
|
)
|
|
|
|
|
2023-07-27 10:03:08 +00:00
|
|
|
func configureHost(ctx context.Context, persistentAuth *auth.PersistentAuth, args []string, argIndex int) error {
|
2023-07-12 15:36:09 +00:00
|
|
|
if len(args) > argIndex {
|
|
|
|
persistentAuth.Host = args[argIndex]
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
host, err := promptForHost(ctx)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
persistentAuth.Host = host
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2023-11-09 16:38:45 +00:00
|
|
|
const minimalDbConnectVersion = "13.1"
|
|
|
|
|
2023-07-27 10:03:08 +00:00
|
|
|
func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
|
|
|
|
cmd := &cobra.Command{
|
|
|
|
Use: "login [HOST]",
|
|
|
|
Short: "Authenticate this machine",
|
|
|
|
}
|
|
|
|
|
|
|
|
var loginTimeout time.Duration
|
|
|
|
var configureCluster bool
|
|
|
|
cmd.Flags().DurationVar(&loginTimeout, "timeout", auth.DefaultTimeout,
|
|
|
|
"Timeout for completing login challenge in the browser")
|
|
|
|
cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
|
|
|
|
"Prompts to configure cluster")
|
|
|
|
|
|
|
|
cmd.RunE = func(cmd *cobra.Command, args []string) error {
|
2023-07-12 15:36:09 +00:00
|
|
|
ctx := cmd.Context()
|
2023-08-02 09:43:42 +00:00
|
|
|
|
|
|
|
var profileName string
|
|
|
|
profileFlag := cmd.Flag("profile")
|
|
|
|
if profileFlag != nil && profileFlag.Value.String() != "" {
|
|
|
|
profileName = profileFlag.Value.String()
|
2023-09-21 12:38:45 +00:00
|
|
|
} else if cmdio.IsInTTY(ctx) {
|
2023-08-02 09:43:42 +00:00
|
|
|
prompt := cmdio.Prompt(ctx)
|
|
|
|
prompt.Label = "Databricks Profile Name"
|
|
|
|
prompt.Default = persistentAuth.ProfileName()
|
|
|
|
prompt.AllowEdit = true
|
|
|
|
profile, err := prompt.Run()
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
profileName = profile
|
|
|
|
}
|
|
|
|
|
2023-10-11 11:12:18 +00:00
|
|
|
err := setHost(ctx, profileName, persistentAuth, args)
|
2023-08-02 09:43:42 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2023-07-12 15:36:09 +00:00
|
|
|
defer persistentAuth.Close()
|
|
|
|
|
|
|
|
// We need the config without the profile before it's used to initialise new workspace client below.
|
|
|
|
// Otherwise it will complain about non existing profile because it was not yet saved.
|
|
|
|
cfg := config.Config{
|
|
|
|
Host: persistentAuth.Host,
|
|
|
|
AuthType: "databricks-cli",
|
|
|
|
}
|
|
|
|
if cfg.IsAccountClient() && persistentAuth.AccountID == "" {
|
|
|
|
accountId, err := promptForAccountID(ctx)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
persistentAuth.AccountID = accountId
|
2023-01-06 15:15:57 +00:00
|
|
|
}
|
2023-07-12 15:36:09 +00:00
|
|
|
cfg.AccountID = persistentAuth.AccountID
|
2023-06-21 15:51:59 +00:00
|
|
|
|
2023-07-12 15:36:09 +00:00
|
|
|
ctx, cancel := context.WithTimeout(ctx, loginTimeout)
|
2023-01-06 15:15:57 +00:00
|
|
|
defer cancel()
|
2023-06-21 10:58:28 +00:00
|
|
|
|
2023-08-02 09:43:42 +00:00
|
|
|
err = persistentAuth.Challenge(ctx)
|
2023-06-21 10:58:28 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
2023-06-21 15:51:59 +00:00
|
|
|
if configureCluster {
|
|
|
|
w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
ctx := cmd.Context()
|
2023-11-09 16:38:45 +00:00
|
|
|
clusterID, err := cfgpickers.AskForCluster(ctx, w,
|
|
|
|
cfgpickers.WithDatabricksConnect(minimalDbConnectVersion))
|
2023-06-21 15:51:59 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2023-11-09 16:38:45 +00:00
|
|
|
cfg.ClusterID = clusterID
|
2023-06-21 15:51:59 +00:00
|
|
|
}
|
2023-06-21 10:58:28 +00:00
|
|
|
|
2023-09-21 12:38:45 +00:00
|
|
|
if profileName != "" {
|
2023-11-10 14:03:57 +00:00
|
|
|
err = databrickscfg.SaveToProfile(ctx, &config.Config{
|
|
|
|
Profile: profileName,
|
|
|
|
Host: cfg.Host,
|
|
|
|
AuthType: cfg.AuthType,
|
|
|
|
AccountID: cfg.AccountID,
|
|
|
|
ClusterID: cfg.ClusterID,
|
|
|
|
})
|
2023-09-21 12:38:45 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName))
|
2023-06-21 10:58:28 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
return nil
|
2023-07-27 10:03:08 +00:00
|
|
|
}
|
2023-06-21 15:51:59 +00:00
|
|
|
|
2023-07-27 10:03:08 +00:00
|
|
|
return cmd
|
2023-01-06 15:15:57 +00:00
|
|
|
}
|
2023-10-11 11:12:18 +00:00
|
|
|
|
|
|
|
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.
|
2023-11-08 14:50:20 +00:00
|
|
|
_, profiles, err := databrickscfg.LoadProfiles(ctx, func(p databrickscfg.Profile) bool {
|
2023-10-11 11:12:18 +00:00
|
|
|
return p.Name == profileName
|
|
|
|
})
|
2023-11-23 09:04:54 +00:00
|
|
|
// Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow.
|
|
|
|
if !errors.Is(err, databrickscfg.ErrNoConfiguration) {
|
2023-10-11 11:12:18 +00:00
|
|
|
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
|
|
|
|
}
|