Skip to content

Commit f9cac44

Browse files
pieternclaude
pietern
and
claude
authored
Add Lakebase Autoscaling support to psql command (#4399)
## Summary - Add support for the new Lakebase Autoscaling API (projects/branches/endpoints) alongside existing Lakebase Provisioned instances - Introduce `--provisioned` and `--autoscaling` flags to filter by product type during interactive selection - Add `--project`, `--branch`, and `--endpoint` flags for direct autoscaling resource targeting - Auto-select branch/endpoint when only one exists for streamlined UX - Skip retries for non-retryable connection errors (authentication failures, missing roles/databases) ## Backward Compatibility The existing interface remains fully compatible: - `databricks psql my-instance` continues to work for Lakebase Provisioned - `databricks psql` (no args) still shows interactive selection, now including both products - All existing flags and behaviors are preserved ## Test plan - New acceptance tests - Manually tested with both provisioned and autoscaling instances --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
1 parent d7010ff commit f9cac44

File tree

29 files changed

+1685
-341
lines changed

29 files changed

+1685
-341
lines changed

‎NEXT_CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
### CLI
88

9+
* Add Lakebase Autoscaling support to `psql` command ([#4399](https://github.com/databricks/cli/pull/4399))
10+
911
### Bundles
1012

1113
* Add missing values to SchemaGrantPrivilege enum ([#4380](https://github.com/databricks/cli/pull/4380))

‎acceptance/cmd/psql/argument-errors/echo-arguments.sh‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#!/bin/bash
2+
#
3+
# This script prints its arguments and exits.
4+
# The test script renames this script to "psql" in order to capture the arguments that the CLI passes to psql command.
5+
#
6+
echo "echo-arguments.sh was called with the following arguments: $@"
7+
echo "PGPASSWORD=${PGPASSWORD}"
8+
echo "PGSSLMODE=${PGSSLMODE}"
9+
exit 0

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

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎acceptance/cmd/psql/argument-errors/output.txt‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
2+
=== Provisioned instance with --project flag should error:
3+
>>> musterr [CLI] psql my-instance --project foo
4+
Error: cannot use --project, --branch, or --endpoint flags with a provisioned instance name
5+
6+
=== Provisioned instance with --branch flag should error:
7+
>>> musterr [CLI] psql my-instance --branch main
8+
Error: cannot use --project, --branch, or --endpoint flags with a provisioned instance name
9+
10+
=== Provisioned instance with --endpoint flag should error:
11+
>>> musterr [CLI] psql my-instance --endpoint primary
12+
Error: cannot use --project, --branch, or --endpoint flags with a provisioned instance name
13+
14+
=== Full path with conflicting --project flag should error:
15+
>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/primary --project other
16+
Error: --project flag conflicts with project in path: other vs my-project
17+
18+
=== Full path with conflicting --branch flag should error:
19+
>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/primary --branch other
20+
Error: --branch flag conflicts with branch in path: other vs main
21+
22+
=== Full path with conflicting --endpoint flag should error:
23+
>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/primary --endpoint other
24+
Error: --endpoint flag conflicts with endpoint in path: other vs primary
25+
26+
=== Partial path (project+branch) with conflicting --branch flag should error:
27+
>>> musterr [CLI] psql projects/my-project/branches/main --branch other
28+
Error: --branch flag conflicts with branch in path: other vs main
29+
30+
=== Partial path (project only) with --branch flag should supplement:
31+
>>> [CLI] psql projects/my-project --branch main --max-retries 0
32+
Project: My Project
33+
Branch: main
34+
Endpoint: primary
35+
Connecting to read-write endpoint...
36+
echo-arguments.sh was called with the following arguments: --host=my-endpoint.postgres.example.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres
37+
PGPASSWORD=postgres-secret-token
38+
PGSSLMODE=require
39+
40+
=== Partial path (project+branch) with --endpoint flag should supplement:
41+
>>> [CLI] psql projects/my-project/branches/dev --endpoint custom --max-retries 0
42+
Project: My Project
43+
Branch: dev
44+
Endpoint: custom
45+
Connecting to read-write endpoint...
46+
echo-arguments.sh was called with the following arguments: --host=custom-endpoint.postgres.example.com --username=[USERNAME] --port=5432 --dbname=databricks_postgres
47+
PGPASSWORD=postgres-secret-token
48+
PGSSLMODE=require
49+
50+
=== Invalid path with just 'projects/' should error:
51+
>>> musterr [CLI] psql projects/
52+
Error: invalid resource path: missing project ID
53+
54+
=== Invalid path with missing branch ID should error:
55+
>>> musterr [CLI] psql projects/my-project/branches/
56+
Error: invalid resource path: missing branch ID
57+
58+
=== Invalid path with missing endpoint ID should error:
59+
>>> musterr [CLI] psql projects/my-project/branches/main/endpoints/
60+
Error: invalid resource path: missing endpoint ID
61+
62+
=== Provisioned flag with --project should error:
63+
>>> musterr [CLI] psql --provisioned --project foo
64+
Error: cannot use --project, --branch, or --endpoint flags with --provisioned
65+
66+
=== Provisioned flag with autoscaling path should error:
67+
>>> musterr [CLI] psql --provisioned projects/my-project/branches/main/endpoints/primary
68+
Error: cannot use --provisioned flag with an autoscaling resource path
69+
70+
=== Autoscaling flag with provisioned instance should error:
71+
>>> musterr [CLI] psql --autoscaling my-instance
72+
Error: cannot use --autoscaling flag with a provisioned instance name
73+
74+
=== Provisioned and autoscaling flags are mutually exclusive:
75+
>>> musterr [CLI] psql --provisioned --autoscaling
76+
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‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
mv echo-arguments.sh psql
2+
3+
export PATH="$(pwd):$PATH"
4+
5+
title "Provisioned instance with --project flag should error:"
6+
trace musterr $CLI psql my-instance --project foo
7+
8+
title "Provisioned instance with --branch flag should error:"
9+
trace musterr $CLI psql my-instance --branch main
10+
11+
title "Provisioned instance with --endpoint flag should error:"
12+
trace musterr $CLI psql my-instance --endpoint primary
13+
14+
title "Full path with conflicting --project flag should error:"
15+
trace musterr $CLI psql projects/my-project/branches/main/endpoints/primary --project other
16+
17+
title "Full path with conflicting --branch flag should error:"
18+
trace musterr $CLI psql projects/my-project/branches/main/endpoints/primary --branch other
19+
20+
title "Full path with conflicting --endpoint flag should error:"
21+
trace musterr $CLI psql projects/my-project/branches/main/endpoints/primary --endpoint other
22+
23+
title "Partial path (project+branch) with conflicting --branch flag should error:"
24+
trace musterr $CLI psql projects/my-project/branches/main --branch other
25+
26+
title "Partial path (project only) with --branch flag should supplement:"
27+
trace $CLI psql projects/my-project --branch main --max-retries 0
28+
29+
title "Partial path (project+branch) with --endpoint flag should supplement:"
30+
trace $CLI psql projects/my-project/branches/dev --endpoint custom --max-retries 0
31+
32+
title "Invalid path with just 'projects/' should error:"
33+
trace musterr $CLI psql projects/
34+
35+
title "Invalid path with missing branch ID should error:"
36+
trace musterr $CLI psql projects/my-project/branches/
37+
38+
title "Invalid path with missing endpoint ID should error:"
39+
trace musterr $CLI psql projects/my-project/branches/main/endpoints/
40+
41+
title "Provisioned flag with --project should error:"
42+
trace musterr $CLI psql --provisioned --project foo
43+
44+
title "Provisioned flag with autoscaling path should error:"
45+
trace musterr $CLI psql --provisioned projects/my-project/branches/main/endpoints/primary
46+
47+
title "Autoscaling flag with provisioned instance should error:"
48+
trace musterr $CLI psql --autoscaling my-instance
49+
50+
title "Provisioned and autoscaling flags are mutually exclusive:"
51+
trace musterr $CLI psql --provisioned --autoscaling

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

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
# This acceptance test is disabled on Windows runners because
2+
# the current argument capturing method does not work on windows-latest GitHub Runner.
3+
GOOS.windows = false
4+
5+
Ignore = [
6+
"psql"
7+
]
8+
9+
# GetProject response
10+
[[Server]]
11+
Pattern = "GET /api/2.0/postgres/projects/my-project"
12+
Response.Body = '''
13+
{
14+
"name": "projects/my-project",
15+
"status": { "display_name": "My Project" }
16+
}
17+
'''
18+
19+
# Provisioned instance lookup
20+
[[Server]]
21+
Pattern = "GET /api/2.0/database/instances/my-instance"
22+
Response.Body = '''
23+
{
24+
"state": "AVAILABLE",
25+
"pg_version": "14",
26+
"read_write_dns": "my-instance.example.com"
27+
}
28+
'''
29+
30+
# GetBranch responses
31+
[[Server]]
32+
Pattern = "GET /api/2.0/postgres/projects/my-project/branches/main"
33+
Response.Body = '''
34+
{
35+
"name": "projects/my-project/branches/main"
36+
}
37+
'''
38+
39+
[[Server]]
40+
Pattern = "GET /api/2.0/postgres/projects/my-project/branches/dev"
41+
Response.Body = '''
42+
{
43+
"name": "projects/my-project/branches/dev"
44+
}
45+
'''
46+
47+
# Full endpoint for supplement tests
48+
[[Server]]
49+
Pattern = "GET /api/2.0/postgres/projects/my-project/branches/main/endpoints/primary"
50+
Response.Body = '''
51+
{
52+
"name": "projects/my-project/branches/main/endpoints/primary",
53+
"status": {
54+
"current_state": "ACTIVE",
55+
"endpoint_type": "ENDPOINT_TYPE_READ_WRITE",
56+
"hosts": {
57+
"host": "my-endpoint.postgres.example.com"
58+
}
59+
}
60+
}
61+
'''
62+
63+
# Partial path supplement: project only -> needs branches list
64+
[[Server]]
65+
Pattern = "GET /api/2.0/postgres/projects/my-project/branches"
66+
Response.Body = '''
67+
{
68+
"branches": [
69+
{
70+
"name": "projects/my-project/branches/main",
71+
"status": { "default": true }
72+
}
73+
]
74+
}
75+
'''
76+
77+
[[Server]]
78+
Pattern = "GET /api/2.0/postgres/projects/my-project/branches/main/endpoints"
79+
Response.Body = '''
80+
{
81+
"endpoints": [
82+
{
83+
"name": "projects/my-project/branches/main/endpoints/primary",
84+
"status": {
85+
"current_state": "ACTIVE",
86+
"endpoint_type": "ENDPOINT_TYPE_READ_WRITE",
87+
"hosts": { "host": "my-endpoint.postgres.example.com" }
88+
}
89+
}
90+
]
91+
}
92+
'''
93+
94+
# Endpoint for supplement test with --endpoint flag
95+
[[Server]]
96+
Pattern = "GET /api/2.0/postgres/projects/my-project/branches/dev/endpoints/custom"
97+
Response.Body = '''
98+
{
99+
"name": "projects/my-project/branches/dev/endpoints/custom",
100+
"status": {
101+
"current_state": "ACTIVE",
102+
"endpoint_type": "ENDPOINT_TYPE_READ_WRITE",
103+
"hosts": {
104+
"host": "custom-endpoint.postgres.example.com"
105+
}
106+
}
107+
}
108+
'''
109+
110+
# Credential generation
111+
[[Server]]
112+
Pattern = "POST /api/2.0/postgres/credentials"
113+
Response.Body = '''
114+
{
115+
"token": "postgres-secret-token",
116+
"expire_time": "2025-01-01T00:00:00Z"
117+
}
118+
'''
119+
120+
[[Server]]
121+
Pattern = "POST /api/2.0/database/credentials"
122+
Response.Body = '''
123+
{
124+
"token": "provisioned-secret-token"
125+
}
126+
'''

‎acceptance/cmd/psql/completions/output.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,7 @@
22
=== Command should show instances names in autocomplete:
33
my-database
44
another-database
5+
projects/my-project
6+
projects/another-project
57
:4
68
Completion ended with directive: ShellCompDirectiveNoFileComp

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,14 @@ Response.Body = '''
88
]
99
}
1010
'''
11+
12+
[[Server]]
13+
Pattern = "GET /api/2.0/postgres/projects"
14+
Response.Body = '''
15+
{
16+
"projects": [
17+
{"name": "projects/my-project"},
18+
{"name": "projects/another-project"}
19+
]
20+
}
21+
'''

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

Lines changed: 22 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,55 +1,46 @@
11

22
=== Command should use default number of retries:
33
>>> musterr [CLI] psql my-database -- --dbname=db1 -p 3000
4-
Connecting to Databricks Database Instance my-database ...
5-
Postgres version: 14
6-
Database instance status: AVAILABLE
7-
Successfully fetched database credentials
8-
Launching psql session to my-database.my-host.com (attempt 1/3)...
4+
Instance: my-database
5+
Connecting to database instance...
96
Simulating connection failure with exit code '2'
10-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
7+
Connection failed with retryable error: connection failed: psql exited with code 2
118
Connection attempt 1/3 failed, retrying in 1s...
12-
Launching psql session to my-database.my-host.com (attempt 2/3)...
9+
Retrying connection (attempt 2/3)...
1310
Simulating connection failure with exit code '2'
14-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
11+
Connection failed with retryable error: connection failed: psql exited with code 2
1512
Connection attempt 2/3 failed, retrying in 2s...
16-
Launching psql session to my-database.my-host.com (attempt 3/3)...
13+
Retrying connection (attempt 3/3)...
1714
Simulating connection failure with exit code '2'
18-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
19-
Error: failed to connect after 3 attempts, last error: connection failed (retryable): psql exited with code 2
15+
Connection failed with retryable error: connection failed: psql exited with code 2
16+
Error: failed to connect after 3 attempts, last error: connection failed: psql exited with code 2
2017

2118
=== Command should use custom number of retries:
2219
>>> musterr [CLI] psql my-database --max-retries 5 -- --dbname=db1 -p 3000
23-
Connecting to Databricks Database Instance my-database ...
24-
Postgres version: 14
25-
Database instance status: AVAILABLE
26-
Successfully fetched database credentials
27-
Launching psql session to my-database.my-host.com (attempt 1/5)...
20+
Instance: my-database
21+
Connecting to database instance...
2822
Simulating connection failure with exit code '2'
29-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
23+
Connection failed with retryable error: connection failed: psql exited with code 2
3024
Connection attempt 1/5 failed, retrying in 1s...
31-
Launching psql session to my-database.my-host.com (attempt 2/5)...
25+
Retrying connection (attempt 2/5)...
3226
Simulating connection failure with exit code '2'
33-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
27+
Connection failed with retryable error: connection failed: psql exited with code 2
3428
Connection attempt 2/5 failed, retrying in 2s...
35-
Launching psql session to my-database.my-host.com (attempt 3/5)...
29+
Retrying connection (attempt 3/5)...
3630
Simulating connection failure with exit code '2'
37-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
31+
Connection failed with retryable error: connection failed: psql exited with code 2
3832
Connection attempt 3/5 failed, retrying in 4s...
39-
Launching psql session to my-database.my-host.com (attempt 4/5)...
33+
Retrying connection (attempt 4/5)...
4034
Simulating connection failure with exit code '2'
41-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
35+
Connection failed with retryable error: connection failed: psql exited with code 2
4236
Connection attempt 4/5 failed, retrying in 8s...
43-
Launching psql session to my-database.my-host.com (attempt 5/5)...
37+
Retrying connection (attempt 5/5)...
4438
Simulating connection failure with exit code '2'
45-
Connection failed with retryable error: connection failed (retryable): psql exited with code 2
46-
Error: failed to connect after 5 attempts, last error: connection failed (retryable): psql exited with code 2
39+
Connection failed with retryable error: connection failed: psql exited with code 2
40+
Error: failed to connect after 5 attempts, last error: connection failed: psql exited with code 2
4741

4842
=== Command should not use retries:
4943
>>> musterr [CLI] psql my-database --max-retries 0 -- --dbname=db1 -p 3000
50-
Connecting to Databricks Database Instance my-database ...
51-
Postgres version: 14
52-
Database instance status: AVAILABLE
53-
Successfully fetched database credentials
54-
Launching psql with connection to my-database.my-host.com...
44+
Instance: my-database
45+
Connecting to database instance...
5546
Simulating connection failure with exit code '2'

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,3 +30,11 @@ Response.Body = '''
3030
"token": "my-secret-token"
3131
}
3232
'''
33+
34+
[[Server]]
35+
Pattern = "GET /api/2.0/postgres/projects"
36+
Response.Body = '''
37+
{
38+
"projects": []
39+
}
40+
'''

0 commit comments

Comments
 (0)