Add Lakebase Autoscaling support to psql command
#4399
+1,747
−341
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Add lakebasev2 package for new Postgres API (projects/branches/endpoints) - Move existing lakebase code to lakebasev1 package - Update psql command to support both APIs: - Old API: databricks psql <database-instance> - New API: databricks psql --project <id> or projects/<id>/... - Add interactive selection with display names - Auto-select when only one branch/endpoint exists Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Refactor lakebase v1 and v2 packages to contain only library logic (fetching instances, connecting) while the interactive selection and prompting logic now lives in cmd/psql. This improves separation of concerns - libraries handle data and connections, commands handle UI. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Move the psql command building, retry logic, and connection handling into libs/lakebase/psql. Both v1 (Database API) and v2 (Postgres API) now use the shared RetryConfig and ConnectOptions structs. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
The first connection attempt now shows the same message as when retries are disabled. Attempt numbers are only shown on the 2nd and subsequent attempts. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Instead of retrieving the workspace client from context inside the Connect functions, require it as a parameter. This makes the dependency explicit. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Rename "Old Database API" to "Lakebase Provisioned" and "New Postgres API" to "Lakebase Autoscaling" throughout the codebase to align with official Databricks documentation. Also fix the auto-selection acceptance test by configuring the mock server to return a single branch and endpoint, and add test coverage for: - All three flags specified (--project --branch --endpoint) - Extra arguments with --project flag Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Reduce verbosity by showing concise connection status (e.g., "Instance: my-database (provisioned)" instead of multiple status lines). Add FormatEndpointType and FormatEndpointState helpers for cleaner display. Consolidate connection functions and simplify control flow. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…tection Positional argument now takes precedence over flags. Flags can supplement missing components in partial paths but conflict with existing components in the path. Also validates resource path structure and rejects malformed paths with clear error messages. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…ection Extract listAllDatabases function to fetch instances and projects in parallel, used by both autocompletion and interactive selection. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…tput Move the project log line from connectAutoscaling to resolveEndpoint and fetch the project to display its human-readable name. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Return the branch ID string directly instead of the full branch object, matching the pattern used by selectEndpointID. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Follows the same pattern as selectBranchID and selectEndpointID. Displays project display names but returns the project ID. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Update resolveEndpoint to: - Call selectProjectID when projectID is not provided - Perform GET calls on branch and endpoint to validate they exist - Use the fetched objects to build resource names consistently Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
The endpoint is now fetched directly in resolveEndpoint using the SDK. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- psql.go: Main command setup, argument parsing, interactive selection - psql_provisioned.go: connectProvisioned function - psql_autoscaling.go: connectAutoscaling, resolveEndpoint, and select* functions Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…selection These flags limit the interactive selection to a single product segment: - --provisioned: Only show/select from provisioned database instances - --autoscaling: Only show/select from autoscaling projects The flags are mutually exclusive with each other and have appropriate conflict detection with other flags and positional arguments. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Add error classification to distinguish between transient and permanent psql connection failures. Errors like "role does not exist" and "authentication failed" are not retried since they require user action. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Standardize error messages across all selection functions: - Combined: "no Lakebase databases found in workspace" - Provisioned: "no Lakebase Provisioned instances found in workspace" - Autoscaling: "no Lakebase Autoscaling projects found in workspace" Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- Simplify help text with grouped examples instead of numbered list - Add link to Lakebase documentation - Group flags into "Product Selection" and "Autoscaling" sections - Fix usage string to show flags before -- separator Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Show instance name for provisioned and project name for autoscaling in the selection ID, instead of numeric indices. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
pietern
requested review from
andrewnester,
anton-107,
denik and
shreyas-goenka
as code owners
pietern
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
pietern
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
Collaborator
|
Commit: 9b9e19b
18 interesting tests: 8 RECOVERED, 5 KNOWN, 5 SKIP
Top 50 slowest tests (at least 2 minutes):
|
anton-107
approved these changes
Jan 30, 2026
denik
approved these changes
Jan 30, 2026
| func FormatEndpointState(state postgres.EndpointStatusState) string { | ||
| switch state { | ||
| case postgres.EndpointStatusStateActive: | ||
| return "" |
Contributor
There was a problem hiding this comment.
why is this empty string and not "active"?
Contributor
Author
There was a problem hiding this comment.
We don't need to inform the user that the endpoint is in the nominal good state.
I'll refactor this to provide a "suffix" with parens to make the logger below simpler.
| endpointType := FormatEndpointType(endpoint.Status.EndpointType) | ||
| stateDesc := FormatEndpointState(state) | ||
| if stateDesc != "" { | ||
| cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s endpoint (%s)...", endpointType, stateDesc)) |
Contributor
There was a problem hiding this comment.
This branch only executes for state == Idle, intentional?
| cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s endpoint...", endpointType)) | ||
| } | ||
|
|
||
| if endpoint.Status.Hosts == nil || endpoint.Status.Hosts.Host == "" { |
Contributor
There was a problem hiding this comment.
Why some errors are returned before Connecting message and some after? What's the difference?
| @@ -0,0 +1,100 @@ | |||
| package lakebasev2 | |||
Contributor
There was a problem hiding this comment.
Why is this package in libs/, it seems tightly coupled to psql command, so it can live somewhere in cmd/?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
--provisionedand--autoscalingflags to filter by product type during interactive selection--project,--branch, and--endpointflags for direct autoscaling resource targetingBackward Compatibility
The existing interface remains fully compatible:
databricks psql my-instancecontinues to work for Lakebase Provisioneddatabricks psql(no args) still shows interactive selection, now including both productsTest plan