Introduce retries to databricks psql command (#3492)

## Why
<!-- Why are these changes needed? Provide the context that the reviewer
might be missing.
For example, were there any decisions behind the change that are not
reflected in the code itself? -->

Sometimes the connection to the database instance might fail, which can
be remediated with a retry. This change makes the command to retry the
connection 3 times by default.

To disable the retries users can set a `max-retries` value to zero:

```
$ databricks psql my-instance --max-retries 0
```

## Tests
<!-- How have you tested the changes? -->
Added a new acceptance test, modified and existing acceptance test; ran
manual tests on Mac and Windows

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: anton-107

NEXT_CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Notable Changes
 
 ### CLI
+* Introduce retries to `databricks psql` command ([#3492](https://github.com/databricks/cli/pull/3492))
 * Add rule files for coding agents working on the CLI code base ([#3245](https://github.com/databricks/cli/pull/3245))
 
 ### Dependency updates

acceptance/cmd/psql/failing-connection/always-fail.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+#
+# This script prints its arguments and exits.
+# The test script renames this script to "psql" in order to capture the arguments that the CLI passes to psql command.
+#
+echo "Simulating connection failure with exit code '2'"
+exit 2

acceptance/cmd/psql/failing-connection/out.test.toml

@@ -0,0 +1,8 @@
+Local = true
+Cloud = false
+
+[GOOS]
+  windows = false
+
+[EnvMatrix]
+  DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]

acceptance/cmd/psql/failing-connection/output.txt

@@ -0,0 +1,61 @@
+
+=== Command should use default number of retries:
+>>> musterr [CLI] psql my-database -- --dbname=db1 -p 3000
+Connecting to Databricks Database Instance my-database ...
+Postgres version: 14
+Database instance status: AVAILABLE
+Successfully fetched database credentials
+Launching psql session to my-database.my-host.com (attempt 1/3)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection attempt 1/3 failed, retrying in 1s...
+Launching psql session to my-database.my-host.com (attempt 2/3)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection attempt 2/3 failed, retrying in 2s...
+Launching psql session to my-database.my-host.com (attempt 3/3)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Error: failed to connect after 3 attempts, last error: connection failed (retryable): psql exited with code 2
+
+Exit code (musterr): 1
+
+=== Command should use custom number of retries:
+>>> musterr [CLI] psql my-database --max-retries 5 -- --dbname=db1 -p 3000
+Connecting to Databricks Database Instance my-database ...
+Postgres version: 14
+Database instance status: AVAILABLE
+Successfully fetched database credentials
+Launching psql session to my-database.my-host.com (attempt 1/5)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection attempt 1/5 failed, retrying in 1s...
+Launching psql session to my-database.my-host.com (attempt 2/5)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection attempt 2/5 failed, retrying in 2s...
+Launching psql session to my-database.my-host.com (attempt 3/5)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection attempt 3/5 failed, retrying in 4s...
+Launching psql session to my-database.my-host.com (attempt 4/5)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection attempt 4/5 failed, retrying in 8s...
+Launching psql session to my-database.my-host.com (attempt 5/5)...
+Simulating connection failure with exit code '2'
+Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Error: failed to connect after 5 attempts, last error: connection failed (retryable): psql exited with code 2
+
+Exit code (musterr): 1
+
+=== Command should not use retries:
+>>> musterr [CLI] psql my-database --max-retries 0 -- --dbname=db1 -p 3000
+Connecting to Databricks Database Instance my-database ...
+Postgres version: 14
+Database instance status: AVAILABLE
+Successfully fetched database credentials
+Launching psql with connection to my-database.my-host.com...
+Simulating connection failure with exit code '2'
+
+Exit code (musterr): 2

acceptance/cmd/psql/failing-connection/script

@@ -0,0 +1,17 @@
+mv always-fail.sh psql
+
+cleanup() {
+  rm psql
+}
+trap cleanup EXIT
+
+export PATH="$(pwd):$PATH"
+
+title "Command should use default number of retries:"
+trace musterr $CLI psql my-database -- --dbname=db1 -p 3000
+
+title "Command should use custom number of retries:"
+trace musterr $CLI psql my-database --max-retries 5 -- --dbname=db1 -p 3000
+
+title "Command should not use retries:"
+trace musterr $CLI psql my-database --max-retries 0 -- --dbname=db1 -p 3000

acceptance/cmd/psql/failing-connection/test.toml

@@ -0,0 +1,32 @@
+# This acceptance test is disabled on Windows runners because
+# the current argument capturing method does not work on windows-latest GitHub Runner.
+#
+# See PR #3228 for documented attempts to fix this issue:
+# https://github.com/databricks/cli/pull/3228
+GOOS.windows = false
+
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/my-database"
+Response.Body = '''
+{
+    "state": "AVAILABLE",
+    "pg_version": "14",
+    "read_write_dns": "my-database.my-host.com"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/database/instances"
+Response.Body = '''
+{
+    "database_instances": []
+}
+'''
+
+[[Server]]
+Pattern = "POST /api/2.0/database/credentials"
+Response.Body = '''
+{
+    "token": "my-secret-token"
+}
+'''

acceptance/cmd/psql/simple/output.txt

@@ -23,7 +23,7 @@ Connecting to Databricks Database Instance my-database ...
 Postgres version: 14
 Database instance status: AVAILABLE
 Successfully fetched database credentials
-Launching psql with connection to my-database.my-host.com...
+Launching psql session to my-database.my-host.com (attempt 1/3)...
 echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres
 PGPASSWORD=my-secret-token
 PGSSLMODE=require
@@ -34,7 +34,7 @@ Connecting to Databricks Database Instance my-database ...
 Postgres version: 14
 Database instance status: AVAILABLE
 Successfully fetched database credentials
-Launching psql with connection to my-database.my-host.com...
+Launching psql session to my-database.my-host.com (attempt 1/3)...
 echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres -c SELECT * FROM my_table --echo-all
 PGPASSWORD=my-secret-token
 PGSSLMODE=require
@@ -45,7 +45,7 @@ Connecting to Databricks Database Instance my-database ...
 Postgres version: 14
 Database instance status: AVAILABLE
 Successfully fetched database credentials
-Launching psql with connection to my-database.my-host.com...
+Launching psql session to my-database.my-host.com (attempt 1/3)...
 echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --port=5432 --dbname=db1
 PGPASSWORD=my-secret-token
 PGSSLMODE=require
@@ -55,7 +55,7 @@ Connecting to Databricks Database Instance my-database ...
 Postgres version: 14
 Database instance status: AVAILABLE
 Successfully fetched database credentials
-Launching psql with connection to my-database.my-host.com...
+Launching psql session to my-database.my-host.com (attempt 1/3)...
 echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --port=5432 -d db2
 PGPASSWORD=my-secret-token
 PGSSLMODE=require
@@ -66,7 +66,7 @@ Connecting to Databricks Database Instance my-database ...
 Postgres version: 14
 Database instance status: AVAILABLE
 Successfully fetched database credentials
-Launching psql with connection to my-database.my-host.com...
+Launching psql session to my-database.my-host.com (attempt 1/3)...
 echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --dbname=db1 -p 3000
 PGPASSWORD=my-secret-token
 PGSSLMODE=require
@@ -76,6 +76,27 @@ Connecting to Databricks Database Instance my-database ...
 Postgres version: 14
 Database instance status: AVAILABLE
 Successfully fetched database credentials
+Launching psql session to my-database.my-host.com (attempt 1/3)...
+echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] -d db2 --port=3001
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
+
+=== Command should not use retries if max-retries is set to 0:
+>>> [CLI] psql my-database --max-retries 0 -- --dbname=db1 -p 3000
+Connecting to Databricks Database Instance my-database ...
+Postgres version: 14
+Database instance status: AVAILABLE
+Successfully fetched database credentials
+Launching psql with connection to my-database.my-host.com...
+echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --dbname=db1 -p 3000
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
+
+>>> [CLI] psql my-database --max-retries 0 -- -d db2 --port=3001
+Connecting to Databricks Database Instance my-database ...
+Postgres version: 14
+Database instance status: AVAILABLE
+Successfully fetched database credentials
 Launching psql with connection to my-database.my-host.com...
 echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] -d db2 --port=3001
 PGPASSWORD=my-secret-token

acceptance/cmd/psql/simple/script

@@ -29,3 +29,7 @@ trace $CLI psql my-database -- -d db2
 title "Command should use the port from extra arguments when specified:"
 trace $CLI psql my-database -- --dbname=db1 -p 3000
 trace $CLI psql my-database -- -d db2 --port=3001
+
+title "Command should not use retries if max-retries is set to 0:"
+trace $CLI psql my-database --max-retries 0 -- --dbname=db1 -p 3000
+trace $CLI psql my-database --max-retries 0 -- -d db2 --port=3001

cmd/psql/psql.go

@@ -3,6 +3,7 @@ package psql
 import (
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
@@ -26,12 +27,17 @@ func newLakebaseConnectCommand() *cobra.Command {
 
 This command requires a psql client to be installed on your machine for the connection to work.
 
+The command includes automatic retry logic for connection failures. You can configure the retry behavior using the flags below.
+
 You can pass additional arguments to psql after a double-dash (--):
   databricks psql my-database -- -c "SELECT * FROM my_table"
   databricks psql my-database -- --echo-all -d "my-db"
 `,
 	}
 
+	// Add retry configuration flag
+	cmd.Flags().Int("max-retries", 3, "Maximum number of connection retry attempts (set to 0 to disable retries)")
+
 	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
@@ -74,7 +80,17 @@ You can pass additional arguments to psql after a double-dash (--):
 		databaseInstanceName := args[0]
 		extraArgs := args[1:]
 
-		return lakebase.Connect(cmd.Context(), databaseInstanceName, extraArgs...)
+		// Read retry configuration from flags
+		maxRetries, _ := cmd.Flags().GetInt("max-retries")
+
+		retryConfig := lakebase.RetryConfig{
+			MaxRetries:    maxRetries,       // Retries are disables when max-retries is 0
+			InitialDelay:  time.Second,      // Fixed initial delay
+			MaxDelay:      10 * time.Second, // Fixed max delay
+			BackoffFactor: 2.0,              // Fixed backoff factor
+		}
+
+		return lakebase.ConnectWithRetryConfig(cmd.Context(), databaseInstanceName, retryConfig, extraArgs...)
 	}
 
 	cmd.ValidArgsFunction = func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {

libs/lakebase/connect.go

@@ -5,16 +5,53 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"os/exec"
 	"strings"
+	"time"
 
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
-	"github.com/databricks/cli/libs/exec"
+	execlib "github.com/databricks/cli/libs/exec"
 	"github.com/databricks/databricks-sdk-go/service/database"
 	"github.com/google/uuid"
 )
 
-func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...string) error {
+// RetryConfig holds configuration for connection retry behavior
+type RetryConfig struct {
+	MaxRetries    int
+	InitialDelay  time.Duration
+	MaxDelay      time.Duration
+	BackoffFactor float64
+}
+
+// attemptConnection launches psql interactively and returns an error if connection fails
+func attemptConnection(ctx context.Context, args, env []string) error {
+	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
+	cmd.Env = env
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	err := cmd.Run()
+	if err != nil {
+		// Check if the error might be due to connection issues
+		// Since we can't capture stderr when running interactively, we check the exit code
+		var exitError *exec.ExitError
+		if errors.As(err, &exitError) {
+			exitCode := exitError.ExitCode()
+
+			// We do not use the Databricks SDK for checking whether the error is retryable because the call in question is not to the API
+			// psql returns exit code 2 for connection failures
+			if exitCode == 2 {
+				return fmt.Errorf("connection failed (retryable): psql exited with code %d", exitCode)
+			}
+		}
+	}
+
+	return err
+}
+
+func ConnectWithRetryConfig(ctx context.Context, databaseInstanceName string, retryConfig RetryConfig, extraArgs ...string) error {
 	cmdio.LogString(ctx, fmt.Sprintf("Connecting to Databricks Database Instance %s ...", databaseInstanceName))
 
 	w := cmdctx.WorkspaceClient(ctx)
@@ -91,11 +128,61 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		"PGSSLMODE=require",
 	)
 
-	cmdio.LogString(ctx, fmt.Sprintf("Launching psql with connection to %s...", db.ReadWriteDns))
+	// If retries are disabled, go directly to interactive session
+	if retryConfig.MaxRetries <= 0 {
+		cmdio.LogString(ctx, fmt.Sprintf("Launching psql with connection to %s...", db.ReadWriteDns))
+		return execlib.Execv(execlib.ExecvOptions{
+			Args: args,
+			Env:  cmdEnv,
+		})
+	}
 
-	// Execute psql command inline
-	return exec.Execv(exec.ExecvOptions{
-		Args: args,
-		Env:  cmdEnv,
-	})
+	// Try launching psql with retry logic
+	maxRetries := retryConfig.MaxRetries
+	delay := retryConfig.InitialDelay
+
+	var lastErr error
+	for attempt := range maxRetries {
+		if attempt > 0 {
+			cmdio.LogString(ctx, fmt.Sprintf("Connection attempt %d/%d failed, retrying in %v...", attempt, maxRetries, delay))
+
+			// Wait with context cancellation support
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(delay):
+			}
+
+			// Exponential backoff
+			delay = time.Duration(float64(delay) * retryConfig.BackoffFactor)
+			if delay > retryConfig.MaxDelay {
+				delay = retryConfig.MaxDelay
+			}
+		}
+
+		cmdio.LogString(ctx, fmt.Sprintf("Launching psql session to %s (attempt %d/%d)...", db.ReadWriteDns, attempt+1, maxRetries))
+
+		// Try to launch psql and capture the exit status
+		err := attemptConnection(ctx, args, cmdEnv)
+		if err == nil {
+			// psql exited normally (user quit)
+			return nil
+		}
+
+		lastErr = err
+
+		// Check if this is a retryable error
+		// We do not use the Databricks SDK for checking whether the error is retryable because the call in question is not to the API
+		if !strings.Contains(err.Error(), "connection failed (retryable)") {
+			// Non-retryable error, fail immediately
+			return err
+		}
+
+		if attempt < maxRetries {
+			cmdio.LogString(ctx, fmt.Sprintf("Connection failed with retryable error: %v", err))
+		}
+	}
+
+	// All retries exhausted
+	return fmt.Errorf("failed to connect after %d attempts, last error: %w", maxRetries, lastErr)
 }