From 1c8cee864ca06fa44cb755ced19ae997b69ecbe5 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Thu, 26 Jun 2025 12:40:54 +0200
Subject: [PATCH 01/23] Introduce "connect" command to connect to Lakebase
 instances

---
 cmd/cmd.go               |  3 ++
 cmd/connect/connect.go   | 56 ++++++++++++++++++++++++++++
 libs/lakebase/connect.go | 80 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 139 insertions(+)
 create mode 100644 cmd/connect/connect.go
 create mode 100644 libs/lakebase/connect.go

diff --git a/cmd/cmd.go b/cmd/cmd.go
index fe458294f2..ab5b9f0fd3 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"strings"
 
+	"github.com/databricks/cli/cmd/connect"
+
 	"github.com/databricks/cli/cmd/account"
 	"github.com/databricks/cli/cmd/api"
 	"github.com/databricks/cli/cmd/auth"
@@ -101,6 +103,7 @@ func New(ctx context.Context) *cobra.Command {
 	cli.AddCommand(api.New())
 	cli.AddCommand(auth.New())
 	cli.AddCommand(bundle.New())
+	cli.AddCommand(connect.New())
 	cli.AddCommand(configure.New())
 	cli.AddCommand(fs.New())
 	cli.AddCommand(labs.New(ctx))
diff --git a/cmd/connect/connect.go b/cmd/connect/connect.go
new file mode 100644
index 0000000000..23edab6cc7
--- /dev/null
+++ b/cmd/connect/connect.go
@@ -0,0 +1,56 @@
+package connect
+
+import (
+	"errors"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/lakebase"
+	"github.com/spf13/cobra"
+)
+
+func New() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "connect",
+		Short:   "Connect allows you to connect to your databases in your Databricks Workspace.",
+		Long:    "Connect allows you to connect to your databases in your Databricks Workspace. You can connect to Lakebase database instances.",
+		GroupID: "development",
+	}
+
+	cmd.AddCommand(newLakebaseConnectCommand())
+
+	return cmd
+}
+
+func newLakebaseConnectCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "lakebase [DATABASE_INSTANCE_NAME]",
+		Short: "Connect to the specified Lakebase database instance",
+		Args:  root.MaximumNArgs(1),
+		Long: `Connect to the specified Lakebase database instance.
+
+You need to have psql client installed on your machine for this connection to work.
+`,
+	}
+
+	// Wrapper for [root.MustWorkspaceClient] that disables loading authentication configuration from a bundle.
+	mustWorkspaceClient := func(cmd *cobra.Command, args []string) error {
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
+	cmd.PreRunE = mustWorkspaceClient
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		var databaseInstanceName string
+		if len(args) > 0 {
+			databaseInstanceName = args[0]
+		}
+
+		if databaseInstanceName == "" {
+			return errors.New("please specify a database instance name: databricks connect lakebase [DATABASE_INSTANCE_NAME]")
+		}
+
+		return lakebase.Connect(cmd.Context(), databaseInstanceName)
+	}
+
+	return cmd
+}
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
new file mode 100644
index 0000000000..f3f5ff5f78
--- /dev/null
+++ b/libs/lakebase/connect.go
@@ -0,0 +1,80 @@
+package lakebase
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/databricks-sdk-go/service/database"
+	"github.com/google/uuid"
+)
+
+func Connect(ctx context.Context, databaseInstanceName string) error {
+	fmt.Printf("Connecting to Databricks Database Instance %s ...\n", databaseInstanceName)
+
+	w := cmdctx.WorkspaceClient(ctx)
+
+	// get user:
+	user, err := w.CurrentUser.Me(ctx)
+	if err != nil {
+		fmt.Println("Error getting current user: ", err)
+		return err
+	}
+
+	// get database:
+	db, err := w.Database.GetDatabaseInstance(ctx, database.GetDatabaseInstanceRequest{
+		Name: databaseInstanceName,
+	})
+	if err != nil {
+		fmt.Println("Error getting Database Instance: ", err)
+		fmt.Println("Does the database instance exist?")
+		return err
+	}
+
+	fmt.Println("Database status: ", db.State)
+	fmt.Println("Database postgres version: ", db.PgVersion)
+
+	// get credentials:
+	cred, err := w.Database.GenerateDatabaseCredential(ctx, database.GenerateDatabaseCredentialRequest{
+		InstanceNames: []string{databaseInstanceName},
+		RequestId:     uuid.NewString(),
+	})
+	if err != nil {
+		fmt.Println("Error getting database credentials: ", err)
+		return err
+	}
+	fmt.Println("Successfully fetched database credentials")
+
+	// Prepare psql command with connection parameters
+	cmd := exec.CommandContext(ctx, "psql",
+		"--host="+db.ReadWriteDns,
+		"--username="+user.UserName,
+		"--dbname=databricks_postgres",
+		"--port=5432",
+	)
+
+	// Set environment variables for psql
+	cmd.Env = append(os.Environ(),
+		"PGPASSWORD="+cred.Token,
+		"PGSSLMODE=require",
+	)
+
+	// Connect stdin, stdout, and stderr to enable interactive session
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	fmt.Printf("Launching psql with connection to %s...\n", db.ReadWriteDns)
+
+	// Execute psql command
+	err = cmd.Run()
+	if err != nil {
+		fmt.Printf("Error running psql: %v\n", err)
+		fmt.Println("Do you have `psql` installed in your path?")
+		return err
+	}
+
+	return nil
+}

From 930dcb9f01816084b3b9750c8e54debbfd815ac3 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 14:08:39 +0200
Subject: [PATCH 02/23] Psql command (#3209)

<!-- Brief summary of your changes that is easy to understand -->

<!-- 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? -->

<!-- How have you tested the changes? -->

<!-- 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. -->

add a local-only acceptance test

add test output files

execute psql command using exec.Execv

add support for by-passing arguments to psql
---
 acceptance/cmd/psql/echo-arguments.sh |  7 ++++
 acceptance/cmd/psql/out.test.toml     |  5 +++
 acceptance/cmd/psql/output.txt        | 26 +++++++++++++
 acceptance/cmd/psql/script            | 13 +++++++
 acceptance/cmd/psql/test.toml         | 17 ++++++++
 cmd/cmd.go                            |  4 +-
 cmd/connect/connect.go                | 56 ---------------------------
 cmd/psql/psql.go                      | 50 ++++++++++++++++++++++++
 libs/lakebase/connect.go              | 47 +++++++++++-----------
 9 files changed, 145 insertions(+), 80 deletions(-)
 create mode 100644 acceptance/cmd/psql/echo-arguments.sh
 create mode 100644 acceptance/cmd/psql/out.test.toml
 create mode 100644 acceptance/cmd/psql/output.txt
 create mode 100644 acceptance/cmd/psql/script
 create mode 100644 acceptance/cmd/psql/test.toml
 delete mode 100644 cmd/connect/connect.go
 create mode 100644 cmd/psql/psql.go

diff --git a/acceptance/cmd/psql/echo-arguments.sh b/acceptance/cmd/psql/echo-arguments.sh
new file mode 100644
index 0000000000..ecbed9f69b
--- /dev/null
+++ b/acceptance/cmd/psql/echo-arguments.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 "echo-arguments.sh was called with the following arguments: $@"
+exit 0
diff --git a/acceptance/cmd/psql/out.test.toml b/acceptance/cmd/psql/out.test.toml
new file mode 100644
index 0000000000..8f3575be7b
--- /dev/null
+++ b/acceptance/cmd/psql/out.test.toml
@@ -0,0 +1,5 @@
+Local = true
+Cloud = false
+
+[EnvMatrix]
+  DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]
diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
new file mode 100644
index 0000000000..358fad9eba
--- /dev/null
+++ b/acceptance/cmd/psql/output.txt
@@ -0,0 +1,26 @@
+
+>>> musterr [CLI] psql my-database
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+Successfully fetched database credentials
+Launching psql with connection to my-database.my-host.com...
+Error: looking up "psql" failed: exec: "psql": executable file not found in $PATH
+
+Exit code (musterr): 1
+
+>>> [CLI] psql my-database
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+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=databricks_postgres --port=5432
+
+>>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all -d my-db
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+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=databricks_postgres --port=5432 -c SELECT * FROM my_table --echo-all -d my-db
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
new file mode 100644
index 0000000000..71e3632242
--- /dev/null
+++ b/acceptance/cmd/psql/script
@@ -0,0 +1,13 @@
+mv echo-arguments.sh psql
+chmod +x psql
+
+cleanup() {
+  rm psql
+}
+trap cleanup EXIT
+
+trace musterr $CLI psql my-database
+
+export PATH="$PATH:$(pwd)"
+trace $CLI psql my-database
+trace $CLI psql my-database -- -c "SELECT * FROM my_table" --echo-all -d "my-db"
diff --git a/acceptance/cmd/psql/test.toml b/acceptance/cmd/psql/test.toml
new file mode 100644
index 0000000000..78873070f8
--- /dev/null
+++ b/acceptance/cmd/psql/test.toml
@@ -0,0 +1,17 @@
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/my-database"
+Response.Body = '''
+{
+    "state": "RUNNING",
+    "pg_version": "14",
+    "read_write_dns": "my-database.my-host.com"
+}
+'''
+
+[[Server]]
+Pattern = "POST /api/2.0/database/credentials"
+Response.Body = '''
+{
+    "token": "my-secret-token"
+}
+'''
diff --git a/cmd/cmd.go b/cmd/cmd.go
index ab5b9f0fd3..ef90ccca6a 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -4,7 +4,7 @@ import (
 	"context"
 	"strings"
 
-	"github.com/databricks/cli/cmd/connect"
+	"github.com/databricks/cli/cmd/psql"
 
 	"github.com/databricks/cli/cmd/account"
 	"github.com/databricks/cli/cmd/api"
@@ -103,7 +103,7 @@ func New(ctx context.Context) *cobra.Command {
 	cli.AddCommand(api.New())
 	cli.AddCommand(auth.New())
 	cli.AddCommand(bundle.New())
-	cli.AddCommand(connect.New())
+	cli.AddCommand(psql.New())
 	cli.AddCommand(configure.New())
 	cli.AddCommand(fs.New())
 	cli.AddCommand(labs.New(ctx))
diff --git a/cmd/connect/connect.go b/cmd/connect/connect.go
deleted file mode 100644
index 23edab6cc7..0000000000
--- a/cmd/connect/connect.go
+++ /dev/null
@@ -1,56 +0,0 @@
-package connect
-
-import (
-	"errors"
-
-	"github.com/databricks/cli/cmd/root"
-	"github.com/databricks/cli/libs/lakebase"
-	"github.com/spf13/cobra"
-)
-
-func New() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:     "connect",
-		Short:   "Connect allows you to connect to your databases in your Databricks Workspace.",
-		Long:    "Connect allows you to connect to your databases in your Databricks Workspace. You can connect to Lakebase database instances.",
-		GroupID: "development",
-	}
-
-	cmd.AddCommand(newLakebaseConnectCommand())
-
-	return cmd
-}
-
-func newLakebaseConnectCommand() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "lakebase [DATABASE_INSTANCE_NAME]",
-		Short: "Connect to the specified Lakebase database instance",
-		Args:  root.MaximumNArgs(1),
-		Long: `Connect to the specified Lakebase database instance.
-
-You need to have psql client installed on your machine for this connection to work.
-`,
-	}
-
-	// Wrapper for [root.MustWorkspaceClient] that disables loading authentication configuration from a bundle.
-	mustWorkspaceClient := func(cmd *cobra.Command, args []string) error {
-		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
-		return root.MustWorkspaceClient(cmd, args)
-	}
-
-	cmd.PreRunE = mustWorkspaceClient
-	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		var databaseInstanceName string
-		if len(args) > 0 {
-			databaseInstanceName = args[0]
-		}
-
-		if databaseInstanceName == "" {
-			return errors.New("please specify a database instance name: databricks connect lakebase [DATABASE_INSTANCE_NAME]")
-		}
-
-		return lakebase.Connect(cmd.Context(), databaseInstanceName)
-	}
-
-	return cmd
-}
diff --git a/cmd/psql/psql.go b/cmd/psql/psql.go
new file mode 100644
index 0000000000..4806b5b4c3
--- /dev/null
+++ b/cmd/psql/psql.go
@@ -0,0 +1,50 @@
+package psql
+
+import (
+	"errors"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/lakebase"
+	"github.com/spf13/cobra"
+)
+
+func New() *cobra.Command {
+	return newLakebaseConnectCommand()
+}
+
+func newLakebaseConnectCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "psql [DATABASE_INSTANCE_NAME] [-- PSQL_ARGS...]",
+		Short:   "Connect to the specified Database Instance",
+		Args:    cobra.MinimumNArgs(0),
+		GroupID: "database",
+		Long: `Connect to the specified Database Instance.
+
+This command requires a psql client to be installed on your machine for the connection to work.
+
+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"
+`,
+	}
+
+	// Wrapper for [root.MustWorkspaceClient] that disables loading authentication configuration from a bundle.
+	mustWorkspaceClient := func(cmd *cobra.Command, args []string) error {
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
+	cmd.PreRunE = mustWorkspaceClient
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		if len(args) == 0 {
+			return errors.New("please specify a database instance name: databricks psql [DATABASE_INSTANCE_NAME]")
+		}
+
+		databaseInstanceName := args[0]
+		extraArgs := args[1:]
+
+		return lakebase.Connect(cmd.Context(), databaseInstanceName, extraArgs...)
+	}
+
+	return cmd
+}
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index f3f5ff5f78..bb36afdc63 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -4,14 +4,14 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"os/exec"
 
 	"github.com/databricks/cli/libs/cmdctx"
+	"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) error {
+func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...string) error {
 	fmt.Printf("Connecting to Databricks Database Instance %s ...\n", databaseInstanceName)
 
 	w := cmdctx.WorkspaceClient(ctx)
@@ -47,34 +47,37 @@ func Connect(ctx context.Context, databaseInstanceName string) error {
 	}
 	fmt.Println("Successfully fetched database credentials")
 
-	// Prepare psql command with connection parameters
-	cmd := exec.CommandContext(ctx, "psql",
-		"--host="+db.ReadWriteDns,
-		"--username="+user.UserName,
+	// Get current working directory
+	dir, err := os.Getwd()
+	if err != nil {
+		fmt.Printf("Error getting working directory: %v\n", err)
+		return err
+	}
+
+	// Prepare command arguments
+	args := []string{
+		"psql",
+		"--host=" + db.ReadWriteDns,
+		"--username=" + user.UserName,
 		"--dbname=databricks_postgres",
 		"--port=5432",
-	)
+	}
+
+	// Append any extra arguments passed through
+	args = append(args, extraArgs...)
 
 	// Set environment variables for psql
-	cmd.Env = append(os.Environ(),
+	cmdEnv := append(os.Environ(),
 		"PGPASSWORD="+cred.Token,
 		"PGSSLMODE=require",
 	)
 
-	// Connect stdin, stdout, and stderr to enable interactive session
-	cmd.Stdin = os.Stdin
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-
 	fmt.Printf("Launching psql with connection to %s...\n", db.ReadWriteDns)
 
-	// Execute psql command
-	err = cmd.Run()
-	if err != nil {
-		fmt.Printf("Error running psql: %v\n", err)
-		fmt.Println("Do you have `psql` installed in your path?")
-		return err
-	}
-
-	return nil
+	// Execute psql command inline
+	return exec.Execv(exec.ExecvOptions{
+		Args: args,
+		Env:  cmdEnv,
+		Dir:  dir,
+	})
 }

From a6124060594c75e361a1a78decac45003b214583 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 15:34:30 +0200
Subject: [PATCH 03/23] fix acceptance test for help

---
 acceptance/help/output.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/acceptance/help/output.txt b/acceptance/help/output.txt
index ca2a59a9d8..e0303855f0 100644
--- a/acceptance/help/output.txt
+++ b/acceptance/help/output.txt
@@ -126,6 +126,7 @@ Clean Rooms
 
 Database
   database                               Database Instances provide access to a database via REST API or direct SQL.
+  psql                                   Connect to the specified Database Instance
 
 Quality Monitor v2
   quality-monitor-v2                     Manage data quality of UC objects (currently support schema).

From 6f21257725831ddf51370540c179daa7c160b089 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 15:57:13 +0200
Subject: [PATCH 04/23] do not use default dbname if it is provided in extra
 args

---
 acceptance/cmd/psql/output.txt | 26 +++++++++++++++++++++++---
 acceptance/cmd/psql/script     | 11 ++++++++++-
 libs/lakebase/connect.go       | 16 +++++++++++++++-
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index 358fad9eba..db31db2f5a 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -1,4 +1,5 @@
 
+=== The command errors out if psql is not found in path:
 >>> musterr [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
 Database status:  RUNNING
@@ -9,18 +10,37 @@ Error: looking up "psql" failed: exec: "psql": executable file not found in $PAT
 
 Exit code (musterr): 1
 
+=== Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
 Database status:  RUNNING
 Database postgres version:  14
 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=databricks_postgres --port=5432
+echo-arguments.sh was called with the following arguments: --host=my-database.my-host.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres
 
->>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all -d my-db
+=== Command should by-pass extra arguments to psql:
+>>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all
 Connecting to Databricks Database Instance my-database ...
 Database status:  RUNNING
 Database postgres version:  14
 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=databricks_postgres --port=5432 -c SELECT * FROM my_table --echo-all -d my-db
+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
+
+=== Command should use the db name from extra arguments when specified:
+>>> [CLI] psql my-database -- --dbname=db1
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+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] --port=5432 --dbname=db1
+
+>>> [CLI] psql my-database -- -d db2
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+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] --port=5432 -d db2
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index 71e3632242..b4ddca361c 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -6,8 +6,17 @@ cleanup() {
 }
 trap cleanup EXIT
 
+title "The command errors out if psql is not found in path:"
 trace musterr $CLI psql my-database
 
 export PATH="$PATH:$(pwd)"
+
+title "Command should use default db name if it is not specified:"
 trace $CLI psql my-database
-trace $CLI psql my-database -- -c "SELECT * FROM my_table" --echo-all -d "my-db"
+
+title "Command should by-pass extra arguments to psql:"
+trace $CLI psql my-database -- -c "SELECT * FROM my_table" --echo-all
+
+title "Command should use the db name from extra arguments when specified:"
+trace $CLI psql my-database -- --dbname=db1
+trace $CLI psql my-database -- -d db2
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index bb36afdc63..b6630c51b2 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/exec"
@@ -54,15 +55,28 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		return err
 	}
 
+	// Check if database name is already specified in extra arguments
+	hasDbName := false
+	for _, arg := range extraArgs {
+		if arg == "-d" || strings.HasPrefix(arg, "--dbname=") {
+			hasDbName = true
+			break
+		}
+	}
+
 	// Prepare command arguments
 	args := []string{
 		"psql",
 		"--host=" + db.ReadWriteDns,
 		"--username=" + user.UserName,
-		"--dbname=databricks_postgres",
 		"--port=5432",
 	}
 
+	// Add default database name only if not specified in extra arguments
+	if !hasDbName {
+		args = append(args, "--dbname=databricks_postgres")
+	}
+
 	// Append any extra arguments passed through
 	args = append(args, extraArgs...)
 

From 84ad8ba73fc821ca2ab6b316439de9bbdfffa62d Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 16:17:07 +0200
Subject: [PATCH 05/23] do not use default port if it is provided in extra args

---
 acceptance/cmd/psql/echo-arguments.sh |  2 ++
 acceptance/cmd/psql/output.txt        | 29 +++++++++++++++++++++++++++
 acceptance/cmd/psql/script            |  4 ++++
 libs/lakebase/connect.go              | 13 +++++++++---
 4 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/acceptance/cmd/psql/echo-arguments.sh b/acceptance/cmd/psql/echo-arguments.sh
index ecbed9f69b..0facd9b697 100644
--- a/acceptance/cmd/psql/echo-arguments.sh
+++ b/acceptance/cmd/psql/echo-arguments.sh
@@ -4,4 +4,6 @@
 # The test script renames this script to "psql" in order to capture the arguments that the CLI passes to psql command
 #
 echo "echo-arguments.sh was called with the following arguments: $@"
+echo "PGPASSWORD=${PGPASSWORD}"
+echo "PGSSLMODE=${PGSSLMODE}"
 exit 0
diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index db31db2f5a..823c94af47 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -18,6 +18,8 @@ Database postgres version:  14
 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] --port=5432 --dbname=databricks_postgres
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
 
 === Command should by-pass extra arguments to psql:
 >>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all
@@ -27,6 +29,8 @@ Database postgres version:  14
 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] --port=5432 --dbname=databricks_postgres -c SELECT * FROM my_table --echo-all
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
 
 === Command should use the db name from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1
@@ -36,6 +40,8 @@ Database postgres version:  14
 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] --port=5432 --dbname=db1
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2
 Connecting to Databricks Database Instance my-database ...
@@ -44,3 +50,26 @@ Database postgres version:  14
 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] --port=5432 -d db2
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
+
+=== Command should use the port from extra arguments when specified:
+>>> [CLI] psql my-database -- --dbname=db1 -p 3000
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+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 -- -d db2 --port=3001
+Connecting to Databricks Database Instance my-database ...
+Database status:  RUNNING
+Database postgres version:  14
+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
+PGSSLMODE=require
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index b4ddca361c..a96a27c708 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -20,3 +20,7 @@ trace $CLI psql my-database -- -c "SELECT * FROM my_table" --echo-all
 title "Command should use the db name from extra arguments when specified:"
 trace $CLI psql my-database -- --dbname=db1
 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
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index b6630c51b2..3598590e88 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -55,12 +55,15 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		return err
 	}
 
-	// Check if database name is already specified in extra arguments
+	// Check if database name and port are already specified in extra arguments
 	hasDbName := false
+	hasPort := false
 	for _, arg := range extraArgs {
 		if arg == "-d" || strings.HasPrefix(arg, "--dbname=") {
 			hasDbName = true
-			break
+		}
+		if arg == "-p" || strings.HasPrefix(arg, "--port=") {
+			hasPort = true
 		}
 	}
 
@@ -69,7 +72,11 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		"psql",
 		"--host=" + db.ReadWriteDns,
 		"--username=" + user.UserName,
-		"--port=5432",
+	}
+
+	// Add default port only if not specified in extra arguments
+	if !hasPort {
+		args = append(args, "--port=5432")
 	}
 
 	// Add default database name only if not specified in extra arguments

From 9e8153de908c0f9970d9bad9349f24aed5fbd40f Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 16:28:25 +0200
Subject: [PATCH 06/23] remove test for condition when psql is not found in
 path

---
 acceptance/cmd/psql/output.txt | 11 -----------
 acceptance/cmd/psql/script     |  3 ---
 2 files changed, 14 deletions(-)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index 823c94af47..f8c726154d 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -1,15 +1,4 @@
 
-=== The command errors out if psql is not found in path:
->>> musterr [CLI] psql my-database
-Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
-Successfully fetched database credentials
-Launching psql with connection to my-database.my-host.com...
-Error: looking up "psql" failed: exec: "psql": executable file not found in $PATH
-
-Exit code (musterr): 1
-
 === Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index a96a27c708..4de241f39f 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -6,9 +6,6 @@ cleanup() {
 }
 trap cleanup EXIT
 
-title "The command errors out if psql is not found in path:"
-trace musterr $CLI psql my-database
-
 export PATH="$PATH:$(pwd)"
 
 title "Command should use default db name if it is not specified:"

From 967236b7bdb5207ffc30f476f02952b4a97b8cff Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 17:03:00 +0200
Subject: [PATCH 07/23] acceptance test should use fake psql command

---
 acceptance/cmd/psql/output.txt | 6 ++++++
 acceptance/cmd/psql/script     | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index f8c726154d..9825dbf43c 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -1,4 +1,10 @@
 
+=== Test should use a fake psql command:
+>>> psql --version
+echo-arguments.sh was called with the following arguments: --version
+PGPASSWORD=
+PGSSLMODE=
+
 === Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index 4de241f39f..5a4870e0ab 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -8,6 +8,9 @@ trap cleanup EXIT
 
 export PATH="$PATH:$(pwd)"
 
+title "Test should use a fake psql command:"
+trace psql --version
+
 title "Command should use default db name if it is not specified:"
 trace $CLI psql my-database
 

From 2e08e8c06cc3bf08c33900d3e2aafa79cfbcff32 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 9 Jul 2025 17:09:50 +0200
Subject: [PATCH 08/23] change the path in the acc script so that fake psql is
 found in path first

---
 acceptance/cmd/psql/script | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index 5a4870e0ab..3f41350996 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -6,7 +6,7 @@ cleanup() {
 }
 trap cleanup EXIT
 
-export PATH="$PATH:$(pwd)"
+export PATH="$(pwd):$PATH"
 
 title "Test should use a fake psql command:"
 trace psql --version

From 2fc6e37ac997828ed29b591f812536f41d889e41 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Fri, 11 Jul 2025 17:16:55 +0200
Subject: [PATCH 09/23] Disable psql test on Windows runners (#3241)

## Changes
<!-- Brief summary of your changes that is easy to understand -->

## 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? -->

## Tests
<!-- How have you tested the changes? -->

<!-- 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. -->
---
 acceptance/cmd/psql/out.test.toml | 3 +++
 acceptance/cmd/psql/test.toml     | 7 +++++++
 2 files changed, 10 insertions(+)

diff --git a/acceptance/cmd/psql/out.test.toml b/acceptance/cmd/psql/out.test.toml
index 8f3575be7b..04844bff9e 100644
--- a/acceptance/cmd/psql/out.test.toml
+++ b/acceptance/cmd/psql/out.test.toml
@@ -1,5 +1,8 @@
 Local = true
 Cloud = false
 
+[GOOS]
+  windows = false
+
 [EnvMatrix]
   DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]
diff --git a/acceptance/cmd/psql/test.toml b/acceptance/cmd/psql/test.toml
index 78873070f8..06de54f72f 100644
--- a/acceptance/cmd/psql/test.toml
+++ b/acceptance/cmd/psql/test.toml
@@ -1,3 +1,10 @@
+# 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 = '''

From 2d8e3cc35c70a6ae9921a5860c0aa1b99d485e43 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Fri, 11 Jul 2025 17:19:40 +0200
Subject: [PATCH 10/23] set echo-arguments.sh as executable

---
 acceptance/cmd/psql/echo-arguments.sh | 4 ++--
 acceptance/cmd/psql/script            | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)
 mode change 100644 => 100755 acceptance/cmd/psql/echo-arguments.sh

diff --git a/acceptance/cmd/psql/echo-arguments.sh b/acceptance/cmd/psql/echo-arguments.sh
old mode 100644
new mode 100755
index 0facd9b697..29d0539e19
--- a/acceptance/cmd/psql/echo-arguments.sh
+++ b/acceptance/cmd/psql/echo-arguments.sh
@@ -1,7 +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
+# 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 "echo-arguments.sh was called with the following arguments: $@"
 echo "PGPASSWORD=${PGPASSWORD}"
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index 3f41350996..0d32ff2487 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -1,5 +1,4 @@
 mv echo-arguments.sh psql
-chmod +x psql
 
 cleanup() {
   rm psql

From 8e870910d0f5b40781a66e9343b7b8ca20e9f109 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Fri, 11 Jul 2025 17:44:00 +0200
Subject: [PATCH 11/23] improve parameters validation + add tests

---
 acceptance/cmd/psql/output.txt | 12 ++++++++++++
 acceptance/cmd/psql/script     |  6 ++++++
 cmd/psql/psql.go               | 14 ++++++++++++--
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index 9825dbf43c..46e185a07d 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -5,6 +5,18 @@ echo-arguments.sh was called with the following arguments: --version
 PGPASSWORD=
 PGSSLMODE=
 
+=== Command should require instance name:
+>>> musterr [CLI] psql
+Error: please specify exactly one database instance name: databricks psql [DATABASE_INSTANCE_NAME]
+
+Exit code (musterr): 1
+
+=== Command should require instance name even if extra arguments are provided:
+>>> musterr [CLI] psql -- --dbname=my-database
+Error: please specify exactly one database instance name: databricks psql [DATABASE_INSTANCE_NAME]
+
+Exit code (musterr): 1
+
 === Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/script
index 0d32ff2487..d662d35420 100644
--- a/acceptance/cmd/psql/script
+++ b/acceptance/cmd/psql/script
@@ -10,6 +10,12 @@ export PATH="$(pwd):$PATH"
 title "Test should use a fake psql command:"
 trace psql --version
 
+title "Command should require instance name:"
+trace musterr $CLI psql
+
+title "Command should require instance name even if extra arguments are provided:"
+trace musterr $CLI psql -- --dbname=my-database
+
 title "Command should use default db name if it is not specified:"
 trace $CLI psql my-database
 
diff --git a/cmd/psql/psql.go b/cmd/psql/psql.go
index 4806b5b4c3..a8ce0dc42f 100644
--- a/cmd/psql/psql.go
+++ b/cmd/psql/psql.go
@@ -36,8 +36,18 @@ You can pass additional arguments to psql after a double-dash (--):
 
 	cmd.PreRunE = mustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
-		if len(args) == 0 {
-			return errors.New("please specify a database instance name: databricks psql [DATABASE_INSTANCE_NAME]")
+		argsLenAtDash := cmd.ArgsLenAtDash()
+
+		// If -- was used, only count args before the dash
+		var argsBeforeDash int
+		if argsLenAtDash >= 0 {
+			argsBeforeDash = argsLenAtDash
+		} else {
+			argsBeforeDash = len(args)
+		}
+
+		if argsBeforeDash != 1 {
+			return errors.New("please specify exactly one database instance name: databricks psql [DATABASE_INSTANCE_NAME]")
 		}
 
 		databaseInstanceName := args[0]

From 4f3284d46b312de397071d39a0735ed09ff9f11a Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Fri, 11 Jul 2025 17:51:54 +0200
Subject: [PATCH 12/23] improve error messaging

---
 libs/lakebase/connect.go | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index 3598590e88..195bd08513 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -20,8 +20,7 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 	// get user:
 	user, err := w.CurrentUser.Me(ctx)
 	if err != nil {
-		fmt.Println("Error getting current user: ", err)
-		return err
+		return fmt.Errorf("error getting current user: %w", err)
 	}
 
 	// get database:
@@ -29,9 +28,7 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		Name: databaseInstanceName,
 	})
 	if err != nil {
-		fmt.Println("Error getting Database Instance: ", err)
-		fmt.Println("Does the database instance exist?")
-		return err
+		return fmt.Errorf("error getting Database Instance. Please confirm that database instance %s exists: %w", databaseInstanceName, err)
 	}
 
 	fmt.Println("Database status: ", db.State)

From abb336d760f553705d8ceb53e5371278bed37696 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Fri, 11 Jul 2025 17:54:28 +0200
Subject: [PATCH 13/23] add NEXT_CHANGELOG.md

---
 NEXT_CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
index f789869b3e..d3c3d63f36 100644
--- a/NEXT_CHANGELOG.md
+++ b/NEXT_CHANGELOG.md
@@ -7,6 +7,7 @@
 ### Dependency updates
 
 ### CLI
+* Added `databricks psql` command to connect to OLTP Databases ([#3128](https://github.com/databricks/cli/pull/3128))
 
 ### Bundles
 

From 92ed62f1d3e212f52bb2aacc4a8e23aa0c0c1d54 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 14 Jul 2025 15:33:59 +0200
Subject: [PATCH 14/23] omit Args field for psql command

---
 cmd/psql/psql.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cmd/psql/psql.go b/cmd/psql/psql.go
index a8ce0dc42f..f80060d462 100644
--- a/cmd/psql/psql.go
+++ b/cmd/psql/psql.go
@@ -16,7 +16,6 @@ func newLakebaseConnectCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:     "psql [DATABASE_INSTANCE_NAME] [-- PSQL_ARGS...]",
 		Short:   "Connect to the specified Database Instance",
-		Args:    cobra.MinimumNArgs(0),
 		GroupID: "database",
 		Long: `Connect to the specified Database Instance.
 

From 4d71f3b4ce7117e72e5aeb74123efc24668960cf Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 14 Jul 2025 15:36:09 +0200
Subject: [PATCH 15/23] omit Args field for psql command

---
 NEXT_CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
index d3c3d63f36..2765c88626 100644
--- a/NEXT_CHANGELOG.md
+++ b/NEXT_CHANGELOG.md
@@ -7,7 +7,7 @@
 ### Dependency updates
 
 ### CLI
-* Added `databricks psql` command to connect to OLTP Databases ([#3128](https://github.com/databricks/cli/pull/3128))
+* Added `databricks psql` command, that allows to connect to OLTP Databases with a single command ([#3128](https://github.com/databricks/cli/pull/3128))
 
 ### Bundles
 

From c8533839ff8f61a251a0439b5c81c387efda7e8b Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 14 Jul 2025 17:19:03 +0200
Subject: [PATCH 16/23] Update NEXT_CHANGELOG.md

Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>
---
 NEXT_CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md
index 2765c88626..b4f8ab0d88 100644
--- a/NEXT_CHANGELOG.md
+++ b/NEXT_CHANGELOG.md
@@ -7,7 +7,7 @@
 ### Dependency updates
 
 ### CLI
-* Added `databricks psql` command, that allows to connect to OLTP Databases with a single command ([#3128](https://github.com/databricks/cli/pull/3128))
+* Add `databricks psql` command to connect to Lakebase with a single command ([#3128](https://github.com/databricks/cli/pull/3128))
 
 ### Bundles
 

From b868d6f50b0c8c44d461dc056f18b7fc0e5dd6d7 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 14 Jul 2025 17:29:45 +0200
Subject: [PATCH 17/23] replace printing errors with returning formatted errors
 from connect.go

---
 libs/lakebase/connect.go | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index 195bd08513..3da6494b52 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -40,16 +40,14 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		RequestId:     uuid.NewString(),
 	})
 	if err != nil {
-		fmt.Println("Error getting database credentials: ", err)
-		return err
+		return fmt.Errorf("error getting database credentials: %w", err)
 	}
 	fmt.Println("Successfully fetched database credentials")
 
 	// Get current working directory
 	dir, err := os.Getwd()
 	if err != nil {
-		fmt.Printf("Error getting working directory: %v\n", err)
-		return err
+		return fmt.Errorf("error getting working directory: %w", err)
 	}
 
 	// Check if database name and port are already specified in extra arguments

From e28dbb09d1e83a75f9ed90a0113c2c97d9031d60 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 14 Jul 2025 18:53:04 +0200
Subject: [PATCH 18/23] use cmdio.LogString to output lines

---
 acceptance/cmd/psql/output.txt | 24 ++++++++++++------------
 libs/lakebase/connect.go       | 11 ++++++-----
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index 46e185a07d..3ff170c7cc 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -20,8 +20,8 @@ Exit code (musterr): 1
 === Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
+Database status: RUNNING
+Database postgres version: 14
 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] --port=5432 --dbname=databricks_postgres
@@ -31,8 +31,8 @@ PGSSLMODE=require
 === Command should by-pass extra arguments to psql:
 >>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all
 Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
+Database status: RUNNING
+Database postgres version: 14
 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] --port=5432 --dbname=databricks_postgres -c SELECT * FROM my_table --echo-all
@@ -42,8 +42,8 @@ PGSSLMODE=require
 === Command should use the db name from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1
 Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
+Database status: RUNNING
+Database postgres version: 14
 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] --port=5432 --dbname=db1
@@ -52,8 +52,8 @@ PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2
 Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
+Database status: RUNNING
+Database postgres version: 14
 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] --port=5432 -d db2
@@ -63,8 +63,8 @@ PGSSLMODE=require
 === Command should use the port from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1 -p 3000
 Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
+Database status: RUNNING
+Database postgres version: 14
 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
@@ -73,8 +73,8 @@ PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2 --port=3001
 Connecting to Databricks Database Instance my-database ...
-Database status:  RUNNING
-Database postgres version:  14
+Database status: RUNNING
+Database postgres version: 14
 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
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index 3da6494b52..2f22759be6 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -7,13 +7,14 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
 	"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 {
-	fmt.Printf("Connecting to Databricks Database Instance %s ...\n", databaseInstanceName)
+	cmdio.LogString(ctx, fmt.Sprintf("Connecting to Databricks Database Instance %s ...", databaseInstanceName))
 
 	w := cmdctx.WorkspaceClient(ctx)
 
@@ -31,8 +32,8 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		return fmt.Errorf("error getting Database Instance. Please confirm that database instance %s exists: %w", databaseInstanceName, err)
 	}
 
-	fmt.Println("Database status: ", db.State)
-	fmt.Println("Database postgres version: ", db.PgVersion)
+	cmdio.LogString(ctx, fmt.Sprintf("Database status: %s", db.State))
+	cmdio.LogString(ctx, fmt.Sprintf("Database postgres version: %s", db.PgVersion))
 
 	// get credentials:
 	cred, err := w.Database.GenerateDatabaseCredential(ctx, database.GenerateDatabaseCredentialRequest{
@@ -42,7 +43,7 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 	if err != nil {
 		return fmt.Errorf("error getting database credentials: %w", err)
 	}
-	fmt.Println("Successfully fetched database credentials")
+	cmdio.LogString(ctx, "Successfully fetched database credentials")
 
 	// Get current working directory
 	dir, err := os.Getwd()
@@ -88,7 +89,7 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		"PGSSLMODE=require",
 	)
 
-	fmt.Printf("Launching psql with connection to %s...\n", db.ReadWriteDns)
+	cmdio.LogString(ctx, fmt.Sprintf("Launching psql with connection to %s...", db.ReadWriteDns))
 
 	// Execute psql command inline
 	return exec.Execv(exec.ExecvOptions{

From d78ff0e1c33d7a3565d82fa23c721d0c4f3fed8f Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Mon, 14 Jul 2025 18:56:00 +0200
Subject: [PATCH 19/23] change output lines for status and postgres version

---
 acceptance/cmd/psql/output.txt | 24 ++++++++++++------------
 libs/lakebase/connect.go       |  4 ++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index 3ff170c7cc..44aaab4a3a 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -20,8 +20,8 @@ Exit code (musterr): 1
 === Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
-Database status: RUNNING
-Database postgres version: 14
+Database instance status: RUNNING
+Postgres version: 14
 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] --port=5432 --dbname=databricks_postgres
@@ -31,8 +31,8 @@ PGSSLMODE=require
 === Command should by-pass extra arguments to psql:
 >>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all
 Connecting to Databricks Database Instance my-database ...
-Database status: RUNNING
-Database postgres version: 14
+Database instance status: RUNNING
+Postgres version: 14
 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] --port=5432 --dbname=databricks_postgres -c SELECT * FROM my_table --echo-all
@@ -42,8 +42,8 @@ PGSSLMODE=require
 === Command should use the db name from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1
 Connecting to Databricks Database Instance my-database ...
-Database status: RUNNING
-Database postgres version: 14
+Database instance status: RUNNING
+Postgres version: 14
 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] --port=5432 --dbname=db1
@@ -52,8 +52,8 @@ PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2
 Connecting to Databricks Database Instance my-database ...
-Database status: RUNNING
-Database postgres version: 14
+Database instance status: RUNNING
+Postgres version: 14
 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] --port=5432 -d db2
@@ -63,8 +63,8 @@ PGSSLMODE=require
 === Command should use the port from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1 -p 3000
 Connecting to Databricks Database Instance my-database ...
-Database status: RUNNING
-Database postgres version: 14
+Database instance status: RUNNING
+Postgres version: 14
 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
@@ -73,8 +73,8 @@ PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2 --port=3001
 Connecting to Databricks Database Instance my-database ...
-Database status: RUNNING
-Database postgres version: 14
+Database instance status: RUNNING
+Postgres version: 14
 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
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index 2f22759be6..45e666539b 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -32,8 +32,8 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		return fmt.Errorf("error getting Database Instance. Please confirm that database instance %s exists: %w", databaseInstanceName, err)
 	}
 
-	cmdio.LogString(ctx, fmt.Sprintf("Database status: %s", db.State))
-	cmdio.LogString(ctx, fmt.Sprintf("Database postgres version: %s", db.PgVersion))
+	cmdio.LogString(ctx, fmt.Sprintf("Database instance status: %s", db.State))
+	cmdio.LogString(ctx, fmt.Sprintf("Postgres version: %s", db.PgVersion))
 
 	// get credentials:
 	cred, err := w.Database.GenerateDatabaseCredential(ctx, database.GenerateDatabaseCredentialRequest{

From 9c0632e0953894aee0344e0679b4d121966cad75 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Tue, 15 Jul 2025 09:55:14 +0200
Subject: [PATCH 20/23] change output lines for status and postgres version,
 fix the spacing

---
 libs/lakebase/connect.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index 45e666539b..013e206954 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -33,7 +33,7 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 	}
 
 	cmdio.LogString(ctx, fmt.Sprintf("Database instance status: %s", db.State))
-	cmdio.LogString(ctx, fmt.Sprintf("Postgres version: %s", db.PgVersion))
+	cmdio.LogString(ctx, "Postgres version: "+db.PgVersion)
 
 	// get credentials:
 	cred, err := w.Database.GenerateDatabaseCredential(ctx, database.GenerateDatabaseCredentialRequest{

From d864462f952e41dcf115e6c4b533050a2f8121ac Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Tue, 15 Jul 2025 16:11:07 +0200
Subject: [PATCH 21/23] show drop-down with instances name in case the instance
 name is not specified

---
 acceptance/cmd/psql/output.txt |  4 ++--
 acceptance/cmd/psql/test.toml  |  8 ++++++++
 cmd/psql/psql.go               | 30 +++++++++++++++++++++++++++++-
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/output.txt
index 44aaab4a3a..a31ae5f40e 100644
--- a/acceptance/cmd/psql/output.txt
+++ b/acceptance/cmd/psql/output.txt
@@ -7,13 +7,13 @@ PGSSLMODE=
 
 === Command should require instance name:
 >>> musterr [CLI] psql
-Error: please specify exactly one database instance name: databricks psql [DATABASE_INSTANCE_NAME]
+Error: could not find any Database instances in the workspace. Please manually specify required argument: DATABASE_INSTANCE_NAME
 
 Exit code (musterr): 1
 
 === Command should require instance name even if extra arguments are provided:
 >>> musterr [CLI] psql -- --dbname=my-database
-Error: please specify exactly one database instance name: databricks psql [DATABASE_INSTANCE_NAME]
+Error: could not find any Database instances in the workspace. Please manually specify required argument: DATABASE_INSTANCE_NAME
 
 Exit code (musterr): 1
 
diff --git a/acceptance/cmd/psql/test.toml b/acceptance/cmd/psql/test.toml
index 06de54f72f..f0c0b32622 100644
--- a/acceptance/cmd/psql/test.toml
+++ b/acceptance/cmd/psql/test.toml
@@ -15,6 +15,14 @@ Response.Body = '''
 }
 '''
 
+[[Server]]
+Pattern = "GET /api/2.0/database/instances"
+Response.Body = '''
+{
+    "database_instances": []
+}
+'''
+
 [[Server]]
 Pattern = "POST /api/2.0/database/credentials"
 Response.Body = '''
diff --git a/cmd/psql/psql.go b/cmd/psql/psql.go
index f80060d462..21f2f63351 100644
--- a/cmd/psql/psql.go
+++ b/cmd/psql/psql.go
@@ -2,6 +2,11 @@ package psql
 
 import (
 	"errors"
+	"fmt"
+
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/databricks-sdk-go/service/database"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/lakebase"
@@ -35,6 +40,8 @@ You can pass additional arguments to psql after a double-dash (--):
 
 	cmd.PreRunE = mustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+		w := cmdctx.WorkspaceClient(ctx)
 		argsLenAtDash := cmd.ArgsLenAtDash()
 
 		// If -- was used, only count args before the dash
@@ -46,7 +53,28 @@ You can pass additional arguments to psql after a double-dash (--):
 		}
 
 		if argsBeforeDash != 1 {
-			return errors.New("please specify exactly one database instance name: databricks psql [DATABASE_INSTANCE_NAME]")
+			promptSpinner := cmdio.Spinner(ctx)
+			promptSpinner <- "No DATABASE_INSTANCE_NAME argument specified. Loading names for Database instances drop-down."
+			instances, err := w.Database.ListDatabaseInstancesAll(ctx, database.ListDatabaseInstancesRequest{})
+			close(promptSpinner)
+			if err != nil {
+				return fmt.Errorf("failed to load names for Database instances drop-down. Please manually specify required argument: DATABASE_INSTANCE_NAME. Original error: %w", err)
+			}
+			if len(instances) == 0 {
+				return errors.New("could not find any Database instances in the workspace. Please manually specify required argument: DATABASE_INSTANCE_NAME")
+			}
+
+			names := make(map[string]string)
+			for _, instance := range instances {
+				names[instance.Name] = instance.Name
+			}
+
+			name, err := cmdio.Select(ctx, names, "")
+			if err != nil {
+				return err
+			}
+
+			args = append([]string{name}, args...)
 		}
 
 		databaseInstanceName := args[0]

From b7e8e33d6a23e55da6e916e93a03b5f37f8a156c Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 16 Jul 2025 11:09:07 +0200
Subject: [PATCH 22/23] add completions for psql command (#3257)

---
 acceptance/cmd/psql/completions/out.test.toml |  5 ++++
 acceptance/cmd/psql/completions/output.txt    |  6 ++++
 acceptance/cmd/psql/completions/script        |  2 ++
 acceptance/cmd/psql/completions/test.toml     | 10 +++++++
 .../cmd/psql/{ => simple}/echo-arguments.sh   |  0
 .../cmd/psql/{ => simple}/out.test.toml       |  0
 acceptance/cmd/psql/{ => simple}/output.txt   |  0
 acceptance/cmd/psql/{ => simple}/script       |  0
 acceptance/cmd/psql/{ => simple}/test.toml    |  0
 cmd/psql/psql.go                              | 29 ++++++++++++++-----
 10 files changed, 45 insertions(+), 7 deletions(-)
 create mode 100644 acceptance/cmd/psql/completions/out.test.toml
 create mode 100644 acceptance/cmd/psql/completions/output.txt
 create mode 100644 acceptance/cmd/psql/completions/script
 create mode 100644 acceptance/cmd/psql/completions/test.toml
 rename acceptance/cmd/psql/{ => simple}/echo-arguments.sh (100%)
 rename acceptance/cmd/psql/{ => simple}/out.test.toml (100%)
 rename acceptance/cmd/psql/{ => simple}/output.txt (100%)
 rename acceptance/cmd/psql/{ => simple}/script (100%)
 rename acceptance/cmd/psql/{ => simple}/test.toml (100%)

diff --git a/acceptance/cmd/psql/completions/out.test.toml b/acceptance/cmd/psql/completions/out.test.toml
new file mode 100644
index 0000000000..8f3575be7b
--- /dev/null
+++ b/acceptance/cmd/psql/completions/out.test.toml
@@ -0,0 +1,5 @@
+Local = true
+Cloud = false
+
+[EnvMatrix]
+  DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]
diff --git a/acceptance/cmd/psql/completions/output.txt b/acceptance/cmd/psql/completions/output.txt
new file mode 100644
index 0000000000..94241d4a31
--- /dev/null
+++ b/acceptance/cmd/psql/completions/output.txt
@@ -0,0 +1,6 @@
+
+=== Command should show instances names in autocomplete:
+my-database
+another-database
+:4
+Completion ended with directive: ShellCompDirectiveNoFileComp
diff --git a/acceptance/cmd/psql/completions/script b/acceptance/cmd/psql/completions/script
new file mode 100644
index 0000000000..b521c17122
--- /dev/null
+++ b/acceptance/cmd/psql/completions/script
@@ -0,0 +1,2 @@
+title "Command should show instances names in autocomplete:\n"
+$CLI __complete psql ""
diff --git a/acceptance/cmd/psql/completions/test.toml b/acceptance/cmd/psql/completions/test.toml
new file mode 100644
index 0000000000..a01fd713ea
--- /dev/null
+++ b/acceptance/cmd/psql/completions/test.toml
@@ -0,0 +1,10 @@
+[[Server]]
+Pattern = "GET /api/2.0/database/instances"
+Response.Body = '''
+{
+    "database_instances": [
+        {"name": "my-database"},
+        {"name": "another-database"}
+    ]
+}
+'''
diff --git a/acceptance/cmd/psql/echo-arguments.sh b/acceptance/cmd/psql/simple/echo-arguments.sh
similarity index 100%
rename from acceptance/cmd/psql/echo-arguments.sh
rename to acceptance/cmd/psql/simple/echo-arguments.sh
diff --git a/acceptance/cmd/psql/out.test.toml b/acceptance/cmd/psql/simple/out.test.toml
similarity index 100%
rename from acceptance/cmd/psql/out.test.toml
rename to acceptance/cmd/psql/simple/out.test.toml
diff --git a/acceptance/cmd/psql/output.txt b/acceptance/cmd/psql/simple/output.txt
similarity index 100%
rename from acceptance/cmd/psql/output.txt
rename to acceptance/cmd/psql/simple/output.txt
diff --git a/acceptance/cmd/psql/script b/acceptance/cmd/psql/simple/script
similarity index 100%
rename from acceptance/cmd/psql/script
rename to acceptance/cmd/psql/simple/script
diff --git a/acceptance/cmd/psql/test.toml b/acceptance/cmd/psql/simple/test.toml
similarity index 100%
rename from acceptance/cmd/psql/test.toml
rename to acceptance/cmd/psql/simple/test.toml
diff --git a/cmd/psql/psql.go b/cmd/psql/psql.go
index 21f2f63351..928784a50b 100644
--- a/cmd/psql/psql.go
+++ b/cmd/psql/psql.go
@@ -32,13 +32,7 @@ You can pass additional arguments to psql after a double-dash (--):
 `,
 	}
 
-	// Wrapper for [root.MustWorkspaceClient] that disables loading authentication configuration from a bundle.
-	mustWorkspaceClient := func(cmd *cobra.Command, args []string) error {
-		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
-		return root.MustWorkspaceClient(cmd, args)
-	}
-
-	cmd.PreRunE = mustWorkspaceClient
+	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 		w := cmdctx.WorkspaceClient(ctx)
@@ -83,5 +77,26 @@ You can pass additional arguments to psql after a double-dash (--):
 		return lakebase.Connect(cmd.Context(), databaseInstanceName, extraArgs...)
 	}
 
+	cmd.ValidArgsFunction = func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		err := root.MustWorkspaceClient(cmd, args)
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		ctx := cmd.Context()
+		w := cmdctx.WorkspaceClient(ctx)
+		instances, err := w.Database.ListDatabaseInstancesAll(ctx, database.ListDatabaseInstancesRequest{})
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		var names []string
+		for _, instance := range instances {
+			names = append(names, instance.Name)
+		}
+
+		return names, cobra.ShellCompDirectiveNoFileComp
+	}
+
 	return cmd
 }

From 99cda31561f1a5c1ffb36eb6cca085f25892b2c7 Mon Sep 17 00:00:00 2001
From: Anton Nekipelov <226657+anton-107@users.noreply.github.com>
Date: Wed, 16 Jul 2025 11:52:54 +0200
Subject: [PATCH 23/23] only try to connect if database instance is in
 AVAILABLE state

---
 .../cmd/psql/not-available/out.test.toml      |  5 +++
 acceptance/cmd/psql/not-available/output.txt  | 44 +++++++++++++++++++
 acceptance/cmd/psql/not-available/script      |  6 +++
 acceptance/cmd/psql/not-available/test.toml   | 44 +++++++++++++++++++
 acceptance/cmd/psql/simple/output.txt         | 12 ++---
 acceptance/cmd/psql/simple/test.toml          |  2 +-
 libs/lakebase/connect.go                      | 10 ++++-
 7 files changed, 115 insertions(+), 8 deletions(-)
 create mode 100644 acceptance/cmd/psql/not-available/out.test.toml
 create mode 100644 acceptance/cmd/psql/not-available/output.txt
 create mode 100644 acceptance/cmd/psql/not-available/script
 create mode 100644 acceptance/cmd/psql/not-available/test.toml

diff --git a/acceptance/cmd/psql/not-available/out.test.toml b/acceptance/cmd/psql/not-available/out.test.toml
new file mode 100644
index 0000000000..8f3575be7b
--- /dev/null
+++ b/acceptance/cmd/psql/not-available/out.test.toml
@@ -0,0 +1,5 @@
+Local = true
+Cloud = false
+
+[EnvMatrix]
+  DATABRICKS_CLI_DEPLOYMENT = ["terraform", "direct-exp"]
diff --git a/acceptance/cmd/psql/not-available/output.txt b/acceptance/cmd/psql/not-available/output.txt
new file mode 100644
index 0000000000..81431542da
--- /dev/null
+++ b/acceptance/cmd/psql/not-available/output.txt
@@ -0,0 +1,44 @@
+
+=== Command should exit if the instance is not available
+>>> musterr [CLI] psql database-starting
+Connecting to Databricks Database Instance database-starting ...
+Postgres version: 14
+Database instance status: STARTING
+Please retry when the instance becomes available
+Error: database instance is not ready for accepting connections
+
+Exit code (musterr): 1
+
+>>> musterr [CLI] psql database-updating
+Connecting to Databricks Database Instance database-updating ...
+Postgres version: 14
+Database instance status: UPDATING
+Please retry when the instance becomes available
+Error: database instance is not ready for accepting connections
+
+Exit code (musterr): 1
+
+>>> musterr [CLI] psql database-stopped
+Connecting to Databricks Database Instance database-stopped ...
+Postgres version: 14
+Database instance status: STOPPED
+Error: database instance is not ready for accepting connections
+
+Exit code (musterr): 1
+
+>>> musterr [CLI] psql database-failing-over
+Connecting to Databricks Database Instance database-failing-over ...
+Postgres version: 14
+Database instance status: FAILING_OVER
+Please retry when the instance becomes available
+Error: database instance is not ready for accepting connections
+
+Exit code (musterr): 1
+
+>>> musterr [CLI] psql database-deleting
+Connecting to Databricks Database Instance database-deleting ...
+Postgres version: 14
+Database instance status: DELETING
+Error: database instance is not ready for accepting connections
+
+Exit code (musterr): 1
diff --git a/acceptance/cmd/psql/not-available/script b/acceptance/cmd/psql/not-available/script
new file mode 100644
index 0000000000..2a45492ab7
--- /dev/null
+++ b/acceptance/cmd/psql/not-available/script
@@ -0,0 +1,6 @@
+title "Command should exit if the instance is not available"
+trace musterr $CLI psql database-starting
+trace musterr $CLI psql database-updating
+trace musterr $CLI psql database-stopped
+trace musterr $CLI psql database-failing-over
+trace musterr $CLI psql database-deleting
diff --git a/acceptance/cmd/psql/not-available/test.toml b/acceptance/cmd/psql/not-available/test.toml
new file mode 100644
index 0000000000..3308ca8a75
--- /dev/null
+++ b/acceptance/cmd/psql/not-available/test.toml
@@ -0,0 +1,44 @@
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/database-starting"
+Response.Body = '''
+{
+    "state": "STARTING",
+    "pg_version": "14"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/database-updating"
+Response.Body = '''
+{
+    "state": "UPDATING",
+    "pg_version": "14"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/database-stopped"
+Response.Body = '''
+{
+    "state": "STOPPED",
+    "pg_version": "14"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/database-failing-over"
+Response.Body = '''
+{
+    "state": "FAILING_OVER",
+    "pg_version": "14"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/database-deleting"
+Response.Body = '''
+{
+    "state": "DELETING",
+    "pg_version": "14"
+}
+'''
diff --git a/acceptance/cmd/psql/simple/output.txt b/acceptance/cmd/psql/simple/output.txt
index a31ae5f40e..9168442b1a 100644
--- a/acceptance/cmd/psql/simple/output.txt
+++ b/acceptance/cmd/psql/simple/output.txt
@@ -20,8 +20,8 @@ Exit code (musterr): 1
 === Command should use default db name if it is not specified:
 >>> [CLI] psql my-database
 Connecting to Databricks Database Instance my-database ...
-Database instance status: RUNNING
 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] --port=5432 --dbname=databricks_postgres
@@ -31,8 +31,8 @@ PGSSLMODE=require
 === Command should by-pass extra arguments to psql:
 >>> [CLI] psql my-database -- -c SELECT * FROM my_table --echo-all
 Connecting to Databricks Database Instance my-database ...
-Database instance status: RUNNING
 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] --port=5432 --dbname=databricks_postgres -c SELECT * FROM my_table --echo-all
@@ -42,8 +42,8 @@ PGSSLMODE=require
 === Command should use the db name from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1
 Connecting to Databricks Database Instance my-database ...
-Database instance status: RUNNING
 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] --port=5432 --dbname=db1
@@ -52,8 +52,8 @@ PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2
 Connecting to Databricks Database Instance my-database ...
-Database instance status: RUNNING
 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] --port=5432 -d db2
@@ -63,8 +63,8 @@ PGSSLMODE=require
 === Command should use the port from extra arguments when specified:
 >>> [CLI] psql my-database -- --dbname=db1 -p 3000
 Connecting to Databricks Database Instance my-database ...
-Database instance status: RUNNING
 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
@@ -73,8 +73,8 @@ PGSSLMODE=require
 
 >>> [CLI] psql my-database -- -d db2 --port=3001
 Connecting to Databricks Database Instance my-database ...
-Database instance status: RUNNING
 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
diff --git a/acceptance/cmd/psql/simple/test.toml b/acceptance/cmd/psql/simple/test.toml
index f0c0b32622..ebf13980a3 100644
--- a/acceptance/cmd/psql/simple/test.toml
+++ b/acceptance/cmd/psql/simple/test.toml
@@ -9,7 +9,7 @@ GOOS.windows = false
 Pattern = "GET /api/2.0/database/instances/my-database"
 Response.Body = '''
 {
-    "state": "RUNNING",
+    "state": "AVAILABLE",
     "pg_version": "14",
     "read_write_dns": "my-database.my-host.com"
 }
diff --git a/libs/lakebase/connect.go b/libs/lakebase/connect.go
index 013e206954..1afacfcea4 100644
--- a/libs/lakebase/connect.go
+++ b/libs/lakebase/connect.go
@@ -2,6 +2,7 @@ package lakebase
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"strings"
@@ -32,8 +33,15 @@ func Connect(ctx context.Context, databaseInstanceName string, extraArgs ...stri
 		return fmt.Errorf("error getting Database Instance. Please confirm that database instance %s exists: %w", databaseInstanceName, err)
 	}
 
-	cmdio.LogString(ctx, fmt.Sprintf("Database instance status: %s", db.State))
 	cmdio.LogString(ctx, "Postgres version: "+db.PgVersion)
+	cmdio.LogString(ctx, fmt.Sprintf("Database instance status: %s", db.State))
+
+	if db.State != database.DatabaseInstanceStateAvailable {
+		if db.State == database.DatabaseInstanceStateStarting || db.State == database.DatabaseInstanceStateUpdating || db.State == database.DatabaseInstanceStateFailingOver {
+			cmdio.LogString(ctx, "Please retry when the instance becomes available")
+		}
+		return errors.New("database instance is not ready for accepting connections")
+	}
 
 	// get credentials:
 	cred, err := w.Database.GenerateDatabaseCredential(ctx, database.GenerateDatabaseCredentialRequest{