Uh oh!
There was an error while loading. Please reload this page.
Add Lakebase Autoscaling support to psql command - #4399
Merged
Conversation
- 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 ownersJanuary 29, 2026 22:15
pieterntemporarily deployed
to
test-trigger-is
January 29, 2026 22:15 — with 
GitHub Actions
Inactive
pieterntemporarily deployed
to
test-trigger-is
January 29, 2026 22:27 — with GitHub Actions
Inactive
Collaborator
Commit: 681a029
19 interesting tests: 7 KNOWN, 5 SKIP, 5 RECOVERED, 2 flaky
Top 50 slowest tests (at least 2 minutes):
|
anton-107
approved these changes
Jan 30, 2026
denik
approved these changes
Jan 30, 2026
Uh oh!
There was an error while loading. Please reload this page.
| 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"?
ContributorAuthor
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?
ContributorAuthor
There was a problem hiding this comment.
Leftover from before. I consolidated and moved the log to right before the psql.Connect call.
| @@ -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/?
ContributorAuthor
There was a problem hiding this comment.
I inlined the functions from these packages into cmd/psql and moved libs/lakebase/psql to libs/psql.
This change consolidates Lakebase-specific connection logic directly into
the psql command, eliminating unnecessary abstraction layers.
Key changes:
- Inline v1 and v2 connection helpers into psql_{provisioned,autoscaling}.go
- Move libs/lakebase/psql to libs/psql (it's a general PostgreSQL client)
- Remove libs/lakebase/v1 and libs/lakebase/v2 packages entirely
- Extract and inline ExtractIDFromName as a local helper in psql_autoscaling.go
- Standardize error messages to use "failed to" pattern consistently
- Add documented constants for default database names
- Fix autoscaling default database from "postgres" to "databricks_postgres"
This simplifies the codebase by having all psql connection logic in one place,
making it easier to understand and maintain. The lakebase libraries were only
used by the psql command, so inlining them reduces indirection without loss
of functionality.
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>pieterntemporarily deployed
to
test-trigger-is
January 30, 2026 15:54 — with GitHub Actions
Inactive
Update expected test outputs to reflect the autoscaling default database change from 'postgres' to 'lakebase'. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Reorder operations so that credential generation and host validation happen before logging "Connecting...". This ensures we only log the connection message after all prerequisite checks have passed. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
pieterntemporarily deployed
to
test-trigger-is
January 30, 2026 16:01 — with GitHub Actions
Inactive
pietern
enabled auto-merge
January 30, 2026 16:02
Uh oh!
There was an error while loading. Please reload this page.
github-merge-queueBot
pushed a commit
that referenced
this pull request
Feb 4, 2026
## Changes This PR adds support for Lakebase Autoscaling resources, enabling declarative configuration of projects, branches, and endpoints. For more information about the product, see https://docs.databricks.com/aws/en/oltp/projects/. These APIs use a "spec" field to capture the user intent in requests and a "status" field to contain the applied values in responses. To match DABs UX patterns, this PR embeds the "spec" type at the top level of the resource. ## Example Example configuration: ```yaml resources: postgres_projects: my_db: project_id: dabs-test-prod-app display_name: "Production Database" pg_version: 17 postgres_branches: main: parent: ${resources.postgres_projects.my_db.id} branch_id: main is_protected: false no_expiry: true postgres_endpoints: primary: parent: ${resources.postgres_branches.main.id} endpoint_id: primary endpoint_type: ENDPOINT_TYPE_READ_WRITE autoscaling_limit_min_cu: 0.5 autoscaling_limit_max_cu: 4 ``` After deploying this, you can access the endpoint using the `psql` command (see #4399): ``` % databricks psql projects/dabs-test-prod-app/branches/main/endpoints/primary -- -c 'select 1;' Project: Production Database Branch: main Endpoint: primary Connecting to read-write endpoint... ?column? ---------- 1 (1 row) ``` **Note:** the default database name assumed by the command is `databricks_postgres`. If this doesn't work, you can try `--dbname lakebase`. ## Tests * Acceptance tests on AWS. * API coverage in the testserver. * Manual testing. # Known limitations * Resource IDs are required: project_id, branch_id, and endpoint_id must be explicitly specified. They're planned to be optional and auto-generated by the server. * ID field updates in Terraform don't trigger recreation as expected. * Endpoint reconciliation conflicts prevent reliable Terraform-based recreate tests for endpoints. --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
deco-sdk-taggingBot
added a commit
that referenced
this pull request
Feb 5, 2026
## Release v0.287.0 ### CLI * Add Lakebase Autoscaling support to `psql` command ([#4399](#4399)) * Add `workspace_id` to the profiles command ([#4404](#4404)) ### Bundles * Add support for Lakebase resources (postgres_projects, postgres_branches, postgres_endpoints) ([#4423](#4423)) * Add missing values to SchemaGrantPrivilege enum ([#4380](#4380)) * Added support for UC catalogs (only in direct mode) ([#4342](#4342)) * engine/direct: Fix updates to fields being ignored for database_catalogs, synced_database_tables (([#4388](#4388))) * engine/direct: Fix migration for database_instances, database_catalog, synced_database_tables ([#4424](#4424)) * Show warning message if valueFrom is used instead of value_from in app config section ([$4429](#4429))
denik pushed a commit
that referenced
this pull request
May 20, 2026
## 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>
denik pushed a commit
that referenced
this pull request
May 20, 2026
## Changes This PR adds support for Lakebase Autoscaling resources, enabling declarative configuration of projects, branches, and endpoints. For more information about the product, see https://docs.databricks.com/aws/en/oltp/projects/. These APIs use a "spec" field to capture the user intent in requests and a "status" field to contain the applied values in responses. To match DABs UX patterns, this PR embeds the "spec" type at the top level of the resource. ## Example Example configuration: ```yaml resources: postgres_projects: my_db: project_id: dabs-test-prod-app display_name: "Production Database" pg_version: 17 postgres_branches: main: parent: ${resources.postgres_projects.my_db.id} branch_id: main is_protected: false no_expiry: true postgres_endpoints: primary: parent: ${resources.postgres_branches.main.id} endpoint_id: primary endpoint_type: ENDPOINT_TYPE_READ_WRITE autoscaling_limit_min_cu: 0.5 autoscaling_limit_max_cu: 4 ``` After deploying this, you can access the endpoint using the `psql` command (see #4399): ``` % databricks psql projects/dabs-test-prod-app/branches/main/endpoints/primary -- -c 'select 1;' Project: Production Database Branch: main Endpoint: primary Connecting to read-write endpoint... ?column? ---------- 1 (1 row) ``` **Note:** the default database name assumed by the command is `databricks_postgres`. If this doesn't work, you can try `--dbname lakebase`. ## Tests * Acceptance tests on AWS. * API coverage in the testserver. * Manual testing. # Known limitations * Resource IDs are required: project_id, branch_id, and endpoint_id must be explicitly specified. They're planned to be optional and auto-generated by the server. * ID field updates in Terraform don't trigger recreation as expected. * Endpoint reconciliation conflicts prevent reliable Terraform-based recreate tests for endpoints. --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
denik pushed a commit
that referenced
this pull request
May 20, 2026
## Release v0.287.0 ### CLI * Add Lakebase Autoscaling support to `psql` command ([#4399](#4399)) * Add `workspace_id` to the profiles command ([#4404](#4404)) ### Bundles * Add support for Lakebase resources (postgres_projects, postgres_branches, postgres_endpoints) ([#4423](#4423)) * Add missing values to SchemaGrantPrivilege enum ([#4380](#4380)) * Added support for UC catalogs (only in direct mode) ([#4342](#4342)) * engine/direct: Fix updates to fields being ignored for database_catalogs, synced_database_tables (([#4388](#4388))) * engine/direct: Fix migration for database_instances, database_catalog, synced_database_tables ([#4424](#4424)) * Show warning message if valueFrom is used instead of value_from in app config section ([$4429](#4429))
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
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
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