Add a flag to configure DB Connect Serverless profiles to databricks auth login. (#3728)

## Changes

This PR extends `databricks auth login` with a new
`--configure-serverless` flag that triggers the flow to configure the
profile for DB Connect Serverless.

The serverless configuration flow simply takes care of configuring the
profile by setting `serverless_compute_id = auto` and deleting the value
of `cluster_id` if any. Though, it is meant to be extended to prompt
users to provide additional properties.

Side note: this PR illustrates the opportunity to consolidate the
profile code management.

## Tests

New acceptance test.

---------

Co-authored-by: Pieter Noordhuis <[email protected]>

Author: renaudhartert-db

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### CLI
 
+* Add the `--configure-serverless` flag to `databricks auth login` to configure Databricks Connect to use serverless.
+
 ### Dependency updates
 
 ### Bundles

acceptance/cmd/auth/login-configure-serverless/out.databrickscfg

@@ -0,0 +1,4 @@
+[DEFAULT]
+host                  = [DATABRICKS_URL]
+serverless_compute_id = auto
+auth_type             = databricks-cli

acceptance/cmd/auth/login-configure-serverless/out.test.toml

@@ -0,0 +1,15 @@
+
+=== Initial profile with cluster_id
+[DEFAULT]
+host = [DATABRICKS_URL]
+cluster_id = existing-cluster-123
+
+=== Run auth login with --configure-serverless
+>>> [CLI] auth login --host [DATABRICKS_URL] --profile DEFAULT --configure-serverless
+Profile DEFAULT was successfully saved
+
+=== Profile after auth login with --configure-serverless
+[DEFAULT]
+host                  = [DATABRICKS_URL]
+serverless_compute_id = auto
+auth_type             = databricks-cli

acceptance/cmd/auth/login-configure-serverless/output.txt

@@ -0,0 +1,24 @@
+sethome "./home"
+
+# Create an initial profile with cluster_id to be cleared out.
+cat > "./home/.databrickscfg" <<EOF
+[DEFAULT]
+host = $DATABRICKS_HOST
+cluster_id = existing-cluster-123
+EOF
+
+title "Initial profile with cluster_id\n"
+cat "./home/.databrickscfg"
+
+# Use a fake browser that performs a GET on the authorization URL
+# and follows the redirect back to localhost.
+export BROWSER="browser.py"
+
+title "Run auth login with --configure-serverless"
+trace $CLI auth login --host $DATABRICKS_HOST --profile DEFAULT --configure-serverless
+
+title "Profile after auth login with --configure-serverless\n"
+cat "./home/.databrickscfg"
+
+# Track the .databrickscfg file that was created to surface changes.
+mv "./home/.databrickscfg" "./out.databrickscfg"

acceptance/cmd/auth/login-configure-serverless/script

@@ -0,0 +1,3 @@
+Ignore = [
+    "home"
+]

acceptance/cmd/auth/login-configure-serverless/test.toml

@@ -97,15 +97,23 @@ depends on the existing profiles you have set in your configuration file
 
 	var loginTimeout time.Duration
 	var configureCluster bool
+	var configureServerless bool
 	cmd.Flags().DurationVar(&loginTimeout, "timeout", defaultTimeout,
 		"Timeout for completing login challenge in the browser")
 	cmd.Flags().BoolVar(&configureCluster, "configure-cluster", false,
 		"Prompts to configure cluster")
+	cmd.Flags().BoolVar(&configureServerless, "configure-serverless", false,
+		"Prompts to configure serverless")
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		profileName := cmd.Flag("profile").Value.String()
 
+		// Cluster and Serverless are mutually exclusive.
+		if configureCluster && configureServerless {
+			return errors.New("please either configure serverless or cluster, not both")
+		}
+
 		// If the user has not specified a profile name, prompt for one.
 		if profileName == "" {
 			var err error
@@ -119,22 +127,16 @@ depends on the existing profiles you have set in your configuration file
 			}
 		}
 
+		// Load parameters from the existing profile if any.
 		existingProfile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler)
 		if err != nil {
 			return err
 		}
-
-		// Set the host and account-id based on the provided arguments and flags.
 		err = setHostAndAccountId(ctx, existingProfile, authArguments, args)
 		if err != nil {
 			return err
 		}
 
-		clusterID := ""
-		if existingProfile != nil {
-			clusterID = existingProfile.ClusterID
-		}
-
 		oauthArgument, err := authArguments.ToOAuthArgument()
 		if err != nil {
 			return err
@@ -155,7 +157,6 @@ depends on the existing profiles you have set in your configuration file
 			Host:      authArguments.Host,
 			AccountID: authArguments.AccountID,
 			AuthType:  "databricks-cli",
-			ClusterID: clusterID,
 		}
 		databricksCfgFile := os.Getenv("DATABRICKS_CONFIG_FILE")
 		if databricksCfgFile != "" {
@@ -169,28 +170,44 @@ depends on the existing profiles you have set in your configuration file
 			return err
 		}
 
-		if configureCluster {
+		switch {
+		case configureCluster:
 			w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
 			if err != nil {
 				return err
 			}
-			ctx := cmd.Context()
 			clusterID, err := cfgpickers.AskForCluster(ctx, w,
 				cfgpickers.WithDatabricksConnect(minimalDbConnectVersion))
 			if err != nil {
 				return err
 			}
 			cfg.ClusterID = clusterID
+		case configureServerless:
+			cfg.ClusterID = ""
+			cfg.ServerlessComputeID = "auto"
+		default:
+			// Respect the existing profile if it exists, even if it has
+			// both cluster and serverless configured. Tools relying on
+			// these fields from the profile will need to handle this case.
+			//
+			// TODO: consider whether we should use this an an opportunity
+			// to clean up the profile under the assumption that serverless
+			// is the preferred option.
+			if existingProfile != nil {
+				cfg.ClusterID = existingProfile.ClusterID
+				cfg.ServerlessComputeID = existingProfile.ServerlessComputeID
+			}
 		}
 
 		if profileName != "" {
 			err = databrickscfg.SaveToProfile(ctx, &config.Config{
-				Profile:    profileName,
-				Host:       cfg.Host,
-				AuthType:   cfg.AuthType,
-				AccountID:  cfg.AccountID,
-				ClusterID:  cfg.ClusterID,
-				ConfigFile: cfg.ConfigFile,
+				Profile:             profileName,
+				Host:                cfg.Host,
+				AuthType:            cfg.AuthType,
+				AccountID:           cfg.AccountID,
+				ClusterID:           cfg.ClusterID,
+				ConfigFile:          cfg.ConfigFile,
+				ServerlessComputeID: cfg.ServerlessComputeID,
 			})
 			if err != nil {
 				return err

cmd/auth/login.go

@@ -79,10 +79,11 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct
 			continue
 		}
 		profile := Profile{
-			Name:      v.Name(),
-			Host:      host,
-			AccountID: all["account_id"],
-			ClusterID: all["cluster_id"],
+			Name:                v.Name(),
+			Host:                host,
+			AccountID:           all["account_id"],
+			ClusterID:           all["cluster_id"],
+			ServerlessComputeID: all["serverless_compute_id"],
 		}
 		if fn(profile) {
 			profiles = append(profiles, profile)

libs/databrickscfg/profile/file.go

@@ -10,10 +10,11 @@ import (
 // It should only be used for prompting and filtering.
 // Use its name to construct a config.Config.
 type Profile struct {
-	Name      string
-	Host      string
-	AccountID string
-	ClusterID string
+	Name                string
+	Host                string
+	AccountID           string
+	ClusterID           string
+	ServerlessComputeID string
 }
 
 func (p Profile) Cloud() string {