Introduce "psql" command to connect to Lakebase instances

NEXT_CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Dependency updates
 
 ### CLI
+* Add `databricks psql` command to connect to Lakebase with a single command ([#3128](https://github.com/databricks/cli/pull/3128))
 
 ### Bundles
 

acceptance/cmd/psql/completions/out.test.toml

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

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

acceptance/cmd/psql/completions/script

@@ -0,0 +1,2 @@
+title "Command should show instances names in autocomplete:\n"
+$CLI __complete psql ""

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"}
+    ]
+}
+'''

acceptance/cmd/psql/not-available/out.test.toml

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

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

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

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"
+}
+'''

acceptance/cmd/psql/simple/echo-arguments.sh

@@ -0,0 +1,9 @@
+#!/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: $@"
+echo "PGPASSWORD=${PGPASSWORD}"
+echo "PGSSLMODE=${PGSSLMODE}"
+exit 0

acceptance/cmd/psql/simple/out.test.toml

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

acceptance/cmd/psql/simple/output.txt

@@ -0,0 +1,82 @@
+
+=== Test should use a fake psql command:
+>>> psql --version
+echo-arguments.sh was called with the following arguments: --version
+PGPASSWORD=
+PGSSLMODE=
+
+=== Command should require instance name:
+>>> musterr [CLI] psql
+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: could not find any Database instances in the workspace. Please manually specify required argument: 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 ...
+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
+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
+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] --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
+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] --port=5432 --dbname=db1
+PGPASSWORD=my-secret-token
+PGSSLMODE=require
+
+>>> [CLI] psql my-database -- -d db2
+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] --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 ...
+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 -- -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
+PGSSLMODE=require

acceptance/cmd/psql/simple/script

@@ -0,0 +1,31 @@
+mv echo-arguments.sh psql
+
+cleanup() {
+  rm psql
+}
+trap cleanup EXIT
+
+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
+
+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
+
+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

acceptance/cmd/psql/simple/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/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).

cmd/cmd.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"strings"
 
+	"github.com/databricks/cli/cmd/psql"
+
 	"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(psql.New())
 	cli.AddCommand(configure.New())
 	cli.AddCommand(fs.New())
 	cli.AddCommand(labs.New(ctx))