Add Lakebase Autoscaling support to psql command

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### CLI
 
+* Add Lakebase Autoscaling support to `psql` command ([#4399](https://github.com/databricks/cli/pull/4399))
+
 ### Bundles
 
 * Add missing values to SchemaGrantPrivilege enum ([#4380](https://github.com/databricks/cli/pull/4380))

acceptance/cmd/psql/argument-errors/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/argument-errors/output.txt

@@ -0,0 +1,76 @@
+
+=== Provisioned instance with --project flag should error:
+>>> musterr [CLI] psql my-instance --project foo
+Error: cannot use --project, --branch, or --endpoint flags with a provisioned instance name
+
+=== Provisioned instance with --branch flag should error:
+>>> musterr [CLI] psql my-instance --branch main
+Error: cannot use --project, --branch, or --endpoint flags with a provisioned instance name
+
+=== Provisioned instance with --endpoint flag should error:
+>>> musterr [CLI] psql my-instance --endpoint primary
+Error: cannot use --project, --branch, or --endpoint flags with a provisioned instance name
+
+=== Full path with conflicting --project flag should error:
+>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/primary --project other
+Error: --project flag conflicts with project in path: other vs my-project
+
+=== Full path with conflicting --branch flag should error:
+>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/primary --branch other
+Error: --branch flag conflicts with branch in path: other vs main
+
+=== Full path with conflicting --endpoint flag should error:
+>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/primary --endpoint other
+Error: --endpoint flag conflicts with endpoint in path: other vs primary
+
+=== Partial path (project+branch) with conflicting --branch flag should error:
+>>> musterr [CLI] psql projects/my-project/branches/main --branch other
+Error: --branch flag conflicts with branch in path: other vs main
+
+=== Partial path (project only) with --branch flag should supplement:
+>>> [CLI] psql projects/my-project --branch main --max-retries 0
+Project: My Project
+Branch: main
+Endpoint: primary
+Connecting to read-write endpoint...
+echo-arguments.sh was called with the following arguments: --host=my-endpoint.postgres.example.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres
+PGPASSWORD=postgres-secret-token
+PGSSLMODE=require
+
+=== Partial path (project+branch) with --endpoint flag should supplement:
+>>> [CLI] psql projects/my-project/branches/dev --endpoint custom --max-retries 0
+Project: My Project
+Branch: dev
+Endpoint: custom
+Connecting to read-write endpoint...
+echo-arguments.sh was called with the following arguments: --host=custom-endpoint.postgres.example.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres
+PGPASSWORD=postgres-secret-token
+PGSSLMODE=require
+
+=== Invalid path with just 'projects/' should error:
+>>> musterr [CLI] psql projects/
+Error: invalid resource path: missing project ID
+
+=== Invalid path with missing branch ID should error:
+>>> musterr [CLI] psql projects/my-project/branches/
+Error: invalid resource path: missing branch ID
+
+=== Invalid path with missing endpoint ID should error:
+>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/
+Error: invalid resource path: missing endpoint ID
+
+=== Provisioned flag with --project should error:
+>>> musterr [CLI] psql --provisioned --project foo
+Error: cannot use --project, --branch, or --endpoint flags with --provisioned
+
+=== Provisioned flag with autoscaling path should error:
+>>> musterr [CLI] psql --provisioned projects/my-project/branches/main/endpoints/primary
+Error: cannot use --provisioned flag with an autoscaling resource path
+
+=== Autoscaling flag with provisioned instance should error:
+>>> musterr [CLI] psql --autoscaling my-instance
+Error: cannot use --autoscaling flag with a provisioned instance name
+
+=== Provisioned and autoscaling flags are mutually exclusive:
+>>> musterr [CLI] psql --provisioned --autoscaling
+Error: if any flags in the group [provisioned autoscaling] are set none of the others can be; [autoscaling provisioned] were all set

acceptance/cmd/psql/argument-errors/script

@@ -0,0 +1,51 @@
+mv echo-arguments.sh psql
+
+export PATH="$(pwd):$PATH"
+
+title "Provisioned instance with --project flag should error:"
+trace musterr $CLI psql my-instance --project foo
+
+title "Provisioned instance with --branch flag should error:"
+trace musterr $CLI psql my-instance --branch main
+
+title "Provisioned instance with --endpoint flag should error:"
+trace musterr $CLI psql my-instance --endpoint primary
+
+title "Full path with conflicting --project flag should error:"
+trace musterr $CLI psql projects/my-project/branches/main/endpoints/primary --project other
+
+title "Full path with conflicting --branch flag should error:"
+trace musterr $CLI psql projects/my-project/branches/main/endpoints/primary --branch other
+
+title "Full path with conflicting --endpoint flag should error:"
+trace musterr $CLI psql projects/my-project/branches/main/endpoints/primary --endpoint other
+
+title "Partial path (project+branch) with conflicting --branch flag should error:"
+trace musterr $CLI psql projects/my-project/branches/main --branch other
+
+title "Partial path (project only) with --branch flag should supplement:"
+trace $CLI psql projects/my-project --branch main --max-retries 0
+
+title "Partial path (project+branch) with --endpoint flag should supplement:"
+trace $CLI psql projects/my-project/branches/dev --endpoint custom --max-retries 0
+
+title "Invalid path with just 'projects/' should error:"
+trace musterr $CLI psql projects/
+
+title "Invalid path with missing branch ID should error:"
+trace musterr $CLI psql projects/my-project/branches/
+
+title "Invalid path with missing endpoint ID should error:"
+trace musterr $CLI psql projects/my-project/branches/main/endpoints/
+
+title "Provisioned flag with --project should error:"
+trace musterr $CLI psql --provisioned --project foo
+
+title "Provisioned flag with autoscaling path should error:"
+trace musterr $CLI psql --provisioned projects/my-project/branches/main/endpoints/primary
+
+title "Autoscaling flag with provisioned instance should error:"
+trace musterr $CLI psql --autoscaling my-instance
+
+title "Provisioned and autoscaling flags are mutually exclusive:"
+trace musterr $CLI psql --provisioned --autoscaling

acceptance/cmd/psql/argument-errors/test.toml

@@ -0,0 +1,126 @@
+# This acceptance test is disabled on Windows runners because
+# the current argument capturing method does not work on windows-latest GitHub Runner.
+GOOS.windows = false
+
+Ignore = [
+    "psql"
+]
+
+# GetProject response
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project"
+Response.Body = '''
+{
+    "name": "projects/my-project",
+    "status": { "display_name": "My Project" }
+}
+'''
+
+# Provisioned instance lookup
+[[Server]]
+Pattern = "GET /api/2.0/database/instances/my-instance"
+Response.Body = '''
+{
+    "state": "AVAILABLE",
+    "pg_version": "14",
+    "read_write_dns": "my-instance.example.com"
+}
+'''
+
+# GetBranch responses
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project/branches/main"
+Response.Body = '''
+{
+    "name": "projects/my-project/branches/main"
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project/branches/dev"
+Response.Body = '''
+{
+    "name": "projects/my-project/branches/dev"
+}
+'''
+
+# Full endpoint for supplement tests
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project/branches/main/endpoints/primary"
+Response.Body = '''
+{
+    "name": "projects/my-project/branches/main/endpoints/primary",
+    "status": {
+        "current_state": "ACTIVE",
+        "endpoint_type": "ENDPOINT_TYPE_READ_WRITE",
+        "hosts": {
+            "host": "my-endpoint.postgres.example.com"
+        }
+    }
+}
+'''
+
+# Partial path supplement: project only -> needs branches list
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project/branches"
+Response.Body = '''
+{
+    "branches": [
+        {
+            "name": "projects/my-project/branches/main",
+            "status": { "default": true }
+        }
+    ]
+}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project/branches/main/endpoints"
+Response.Body = '''
+{
+    "endpoints": [
+        {
+            "name": "projects/my-project/branches/main/endpoints/primary",
+            "status": {
+                "current_state": "ACTIVE",
+                "endpoint_type": "ENDPOINT_TYPE_READ_WRITE",
+                "hosts": { "host": "my-endpoint.postgres.example.com" }
+            }
+        }
+    ]
+}
+'''
+
+# Endpoint for supplement test with --endpoint flag
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects/my-project/branches/dev/endpoints/custom"
+Response.Body = '''
+{
+    "name": "projects/my-project/branches/dev/endpoints/custom",
+    "status": {
+        "current_state": "ACTIVE",
+        "endpoint_type": "ENDPOINT_TYPE_READ_WRITE",
+        "hosts": {
+            "host": "custom-endpoint.postgres.example.com"
+        }
+    }
+}
+'''
+
+# Credential generation
+[[Server]]
+Pattern = "POST /api/2.0/postgres/credentials"
+Response.Body = '''
+{
+    "token": "postgres-secret-token",
+    "expire_time": "2025-01-01T00:00:00Z"
+}
+'''
+
+[[Server]]
+Pattern = "POST /api/2.0/database/credentials"
+Response.Body = '''
+{
+    "token": "provisioned-secret-token"
+}
+'''

acceptance/cmd/psql/completions/output.txt

@@ -2,5 +2,7 @@
 === Command should show instances names in autocomplete:
 my-database
 another-database
+projects/my-project
+projects/another-project
 :4
 Completion ended with directive: ShellCompDirectiveNoFileComp

acceptance/cmd/psql/completions/test.toml

@@ -8,3 +8,14 @@ Response.Body = '''
     ]
 }
 '''
+
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects"
+Response.Body = '''
+{
+    "projects": [
+        {"name": "projects/my-project"},
+        {"name": "projects/another-project"}
+    ]
+}
+'''

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

@@ -1,55 +1,46 @@
 
 === 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)...
+Instance: my-database
+Connecting to database instance...
 Simulating connection failure with exit code '2'
-Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection failed with retryable error: connection failed: 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)...
+Retrying connection (attempt 2/3)...
 Simulating connection failure with exit code '2'
-Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection failed with retryable error: connection failed: 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)...
+Retrying connection (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
+Connection failed with retryable error: connection failed: psql exited with code 2
+Error: failed to connect after 3 attempts, last error: connection failed: psql exited with code 2
 
 === 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)...
+Instance: my-database
+Connecting to database instance...
 Simulating connection failure with exit code '2'
-Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection failed with retryable error: connection failed: 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)...
+Retrying connection (attempt 2/5)...
 Simulating connection failure with exit code '2'
-Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection failed with retryable error: connection failed: 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)...
+Retrying connection (attempt 3/5)...
 Simulating connection failure with exit code '2'
-Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection failed with retryable error: connection failed: 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)...
+Retrying connection (attempt 4/5)...
 Simulating connection failure with exit code '2'
-Connection failed with retryable error: connection failed (retryable): psql exited with code 2
+Connection failed with retryable error: connection failed: 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)...
+Retrying connection (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
+Connection failed with retryable error: connection failed: psql exited with code 2
+Error: failed to connect after 5 attempts, last error: connection failed: psql exited with code 2
 
 === 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...
+Instance: my-database
+Connecting to database instance...
 Simulating connection failure with exit code '2'

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

@@ -30,3 +30,11 @@ Response.Body = '''
     "token": "my-secret-token"
 }
 '''
+
+[[Server]]
+Pattern = "GET /api/2.0/postgres/projects"
+Response.Body = '''
+{
+    "projects": []
+}
+'''