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"
|
2024-04-18 11:55:42 +00:00
|
|
|
"runtime"
|
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 {
|
2024-04-18 11:55:42 +00:00
|
|
|
defaultConfigPath := "~/.databrickscfg"
|
|
|
|
if runtime.GOOS == "windows" {
|
|
|
|
defaultConfigPath = "%USERPROFILE%\\.databrickscfg"
|
|
|
|
}
|
2023-07-27 10:03:08 +00:00
|
|
|
cmd := &cobra.Command{
|
|
|
|
Use: "login [HOST]",
|
2024-04-18 11:55:42 +00:00
|
|
|
Short: "Log into a Databricks workspace or account",
|
|
|
|
Long: fmt.Sprintf(`Log into a Databricks workspace or account.
|
|
|
|
This command logs you into the Databricks workspace or account and saves
|
|
|
|
the authentication configuration in a profile (in %s by default).
|
|
|
|
|
|
|
|
This profile can then be used to authenticate other Databricks CLI commands by
|
|
|
|
specifying the --profile flag. This profile can also be used to authenticate
|
|
|
|
other Databricks tooling that supports the Databricks Unified Authentication
|
|
|
|
Specification. This includes the Databricks Go, Python, and Java SDKs. For more information,
|
|
|
|
you can refer to the documentation linked below.
|
|
|
|
AWS: https://docs.databricks.com/dev-tools/auth/index.html
|
|
|
|
Azure: https://learn.microsoft.com/azure/databricks/dev-tools/auth
|
|
|
|
GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html
|
|
|
|
|
|
|
|
|
|
|
|
This command requires a Databricks Host URL (using --host or as a positional argument
|
|
|
|
or implicitly inferred from the specified profile name)
|
|
|
|
and a profile name (using --profile) to be specified. If you don't specify these
|
|
|
|
values, you'll be prompted for values at runtime.
|
|
|
|
|
|
|
|
While this command always logs you into the specified host, the runtime behaviour
|
|
|
|
depends on the existing profiles you have set in your configuration file
|
|
|
|
(at %s by default).
|
|
|
|
|
|
|
|
1. If a profile with the specified name exists and specifies a host, you'll
|
|
|
|
be logged into the host specified by the profile. The profile will be updated
|
|
|
|
to use "databricks-cli" as the auth type if that was not the case before.
|
|
|
|
|
|
|
|
2. If a profile with the specified name exists but does not specify a host,
|
|
|
|
you'll be prompted to specify a host. The profile will be updated to use the
|
|
|
|
specified host. The auth type will be updated to "databricks-cli" if that was
|
|
|
|
not the case before.
|
|
|
|
|
|
|
|
3. If a profile with the specified name exists and specifies a host, but you
|
|
|
|
specify a host using --host (or as the [HOST] positional arg), the profile will
|
|
|
|
be updated to use the newly specified host. The auth type will be updated to
|
|
|
|
"databricks-cli" if that was not the case before.
|
|
|
|
|
|
|
|
4. If a profile with the specified name does not exist, a new profile will be
|
|
|
|
created with the specified host. The auth type will be set to "databricks-cli".
|
|
|
|
`, defaultConfigPath, defaultConfigPath),
|
2023-07-27 10:03:08 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
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.
|
2023-12-04 15:40:52 +00:00
|
|
|
if err != nil && !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
|
|
|
|
}
|