mirror of https://github.com/databricks/cli.git
100 lines
2.7 KiB
Go
100 lines
2.7 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/databricks/cli/libs/auth"
|
|
"github.com/databricks/cli/libs/cmdio"
|
|
"github.com/databricks/cli/libs/databrickscfg"
|
|
"github.com/databricks/databricks-sdk-go"
|
|
"github.com/databricks/databricks-sdk-go/config"
|
|
"github.com/databricks/databricks-sdk-go/service/compute"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var loginTimeout time.Duration
|
|
var configureCluster bool
|
|
|
|
var loginCmd = &cobra.Command{
|
|
Use: "login [HOST]",
|
|
Short: "Authenticate this machine",
|
|
RunE: func(cmd *cobra.Command, args []string) error {
|
|
if perisistentAuth.Host == "" && len(args) == 1 {
|
|
perisistentAuth.Host = args[0]
|
|
}
|
|
|
|
defer perisistentAuth.Close()
|
|
ctx, cancel := context.WithTimeout(cmd.Context(), loginTimeout)
|
|
defer cancel()
|
|
|
|
var profileName string
|
|
profileFlag := cmd.Flag("profile")
|
|
if profileFlag != nil && profileFlag.Value.String() != "" {
|
|
profileName = profileFlag.Value.String()
|
|
} else {
|
|
prompt := cmdio.Prompt(ctx)
|
|
prompt.Label = "Databricks Profile Name"
|
|
prompt.Default = perisistentAuth.ProfileName()
|
|
prompt.AllowEdit = true
|
|
profile, err := prompt.Run()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
profileName = profile
|
|
}
|
|
err := perisistentAuth.Challenge(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// 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: perisistentAuth.Host,
|
|
AccountID: perisistentAuth.AccountID,
|
|
AuthType: "databricks-cli",
|
|
}
|
|
|
|
if configureCluster {
|
|
w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
ctx := cmd.Context()
|
|
|
|
promptSpinner := cmdio.Spinner(ctx)
|
|
promptSpinner <- "Loading list of clusters to select from"
|
|
names, err := w.Clusters.ClusterDetailsClusterNameToClusterIdMap(ctx, compute.ListClustersRequest{})
|
|
close(promptSpinner)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to load clusters list. Original error: %w", err)
|
|
}
|
|
clusterId, err := cmdio.Select(ctx, names, "Choose cluster")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cfg.ClusterID = clusterId
|
|
}
|
|
|
|
cfg.Profile = profileName
|
|
err = databrickscfg.SaveToProfile(ctx, &cfg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName))
|
|
return nil
|
|
},
|
|
}
|
|
|
|
func init() {
|
|
authCmd.AddCommand(loginCmd)
|
|
loginCmd.Flags().DurationVar(&loginTimeout, "timeout", auth.DefaultTimeout,
|
|
"Timeout for completing login challenge in the browser")
|
|
|
|
loginCmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
|
|
"Prompts to configure cluster")
|
|
}
|