From 627af735726fd517301789105c53d864e269b264 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 27 Feb 2026 13:29:48 +0100 Subject: [PATCH 1/4] Respect user's default warehouse override in aitools get-default-warehouse The aitools get-default-warehouse command now checks the user's default warehouse override (set via the SQL UI or CLI) before falling back to server-side default detection. Only CUSTOM overrides are used; LAST_SELECTED is skipped since it requires UI state. New priority order: 1. DATABRICKS_WAREHOUSE_ID env var 2. User's default warehouse override (CUSTOM type) 3. Server-side "default" warehouse 4. First usable warehouse by state --- .../aitools/lib/middlewares/warehouse.go | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/experimental/aitools/lib/middlewares/warehouse.go b/experimental/aitools/lib/middlewares/warehouse.go index 78a62d793a..7cf884e749 100644 --- a/experimental/aitools/lib/middlewares/warehouse.go +++ b/experimental/aitools/lib/middlewares/warehouse.go @@ -101,5 +101,23 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) { }, nil } + // Check user's default warehouse override (set via the SQL UI or CLI). + // Only CUSTOM overrides are used; LAST_SELECTED requires UI state we don't have. + override, err := w.Warehouses.GetDefaultWarehouseOverride(ctx, sql.GetDefaultWarehouseOverrideRequest{ + Name: "default-warehouse-overrides/me", + }) + if err == nil && override.Type == sql.DefaultWarehouseOverrideTypeCustom && override.WarehouseId != "" { + warehouse, err := w.Warehouses.Get(ctx, sql.GetWarehouseRequest{ + Id: override.WarehouseId, + }) + if err == nil { + return &sql.EndpointInfo{ + Id: warehouse.Id, + Name: warehouse.Name, + State: warehouse.State, + }, nil + } + } + return cfgpickers.GetDefaultWarehouse(ctx, w) } From dca62d881e0a9e0b68cacb2f5a7d673c9e1de058 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 27 Feb 2026 13:38:39 +0100 Subject: [PATCH 2/4] Skip deleted/deleting warehouses from override fallback If the user's overridden warehouse is in DELETED or DELETING state, fall through to default detection instead of returning an unusable warehouse. --- experimental/aitools/lib/middlewares/warehouse.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/aitools/lib/middlewares/warehouse.go b/experimental/aitools/lib/middlewares/warehouse.go index 7cf884e749..467e3b842c 100644 --- a/experimental/aitools/lib/middlewares/warehouse.go +++ b/experimental/aitools/lib/middlewares/warehouse.go @@ -110,7 +110,7 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) { warehouse, err := w.Warehouses.Get(ctx, sql.GetWarehouseRequest{ Id: override.WarehouseId, }) - if err == nil { + if err == nil && warehouse.State != sql.StateDeleted && warehouse.State != sql.StateDeleting { return &sql.EndpointInfo{ Id: warehouse.Id, Name: warehouse.Name, From 897f10bf7b4ced2fd516fd427df1ca7a3f809a9d Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 27 Feb 2026 13:43:53 +0100 Subject: [PATCH 3/4] Auto-start stopped warehouses before returning them If the resolved warehouse is in STOPPED or STOPPING state, start it and wait for it to reach RUNNING before returning. This applies to all resolution paths (env var, user override, auto-detection). Extracts warehouse resolution into resolveWarehouse() for clarity. --- .../aitools/lib/middlewares/warehouse.go | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/experimental/aitools/lib/middlewares/warehouse.go b/experimental/aitools/lib/middlewares/warehouse.go index 467e3b842c..dd1ab90db6 100644 --- a/experimental/aitools/lib/middlewares/warehouse.go +++ b/experimental/aitools/lib/middlewares/warehouse.go @@ -8,6 +8,7 @@ import ( "github.com/databricks/cli/experimental/aitools/lib/session" "github.com/databricks/cli/libs/databrickscfg/cfgpickers" "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/sql" ) @@ -85,6 +86,32 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) { return nil, fmt.Errorf("get databricks client: %w", err) } + warehouse, err := resolveWarehouse(ctx, w) + if err != nil { + return nil, err + } + + // Start the warehouse if it's not running. + if warehouse.State == sql.StateStopped || warehouse.State == sql.StateStopping { + wait, err := w.Warehouses.Start(ctx, sql.StartRequest{Id: warehouse.Id}) + if err != nil { + return nil, fmt.Errorf("start warehouse %s: %w", warehouse.Id, err) + } + resp, err := wait.Get() + if err != nil { + return nil, fmt.Errorf("wait for warehouse %s to start: %w", warehouse.Id, err) + } + warehouse.State = resp.State + } + + return warehouse, nil +} + +// resolveWarehouse selects a warehouse using the following priority: +// 1. DATABRICKS_WAREHOUSE_ID env var +// 2. User's default warehouse override (CUSTOM type only) +// 3. Server-side default / first usable warehouse by state +func resolveWarehouse(ctx context.Context, w *databricks.WorkspaceClient) (*sql.EndpointInfo, error) { // first resolve DATABRICKS_WAREHOUSE_ID env variable warehouseID := env.Get(ctx, "DATABRICKS_WAREHOUSE_ID") if warehouseID != "" { From a951b86c21a0c0a49570ead9d84f6ff033889e98 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 27 Feb 2026 13:56:23 +0100 Subject: [PATCH 4/4] Add opt-in auto-start for stopped warehouses GetWarehouseEndpoint and GetWarehouseID now accept an autoStart parameter. When true, a stopped warehouse is started and the call blocks until it reaches RUNNING state. Enabled for query and discover-schema (which need a running warehouse). Disabled for get-default-warehouse and discover (which only report warehouse info). --- experimental/aitools/cmd/discover_schema.go | 2 +- .../aitools/cmd/get_default_warehouse.go | 2 +- experimental/aitools/cmd/query.go | 2 +- .../aitools/lib/middlewares/warehouse.go | 57 ++++++++++++------- .../lib/providers/clitools/discover.go | 2 +- 5 files changed, 40 insertions(+), 25 deletions(-) diff --git a/experimental/aitools/cmd/discover_schema.go b/experimental/aitools/cmd/discover_schema.go index e611764830..9d30e367c0 100644 --- a/experimental/aitools/cmd/discover_schema.go +++ b/experimental/aitools/cmd/discover_schema.go @@ -50,7 +50,7 @@ For each table, returns: sess.Set(middlewares.DatabricksClientKey, w) ctx = session.WithSession(ctx, sess) - warehouseID, err := middlewares.GetWarehouseID(ctx) + warehouseID, err := middlewares.GetWarehouseID(ctx, true) if err != nil { return err } diff --git a/experimental/aitools/cmd/get_default_warehouse.go b/experimental/aitools/cmd/get_default_warehouse.go index 010539d5de..7c8da37b98 100644 --- a/experimental/aitools/cmd/get_default_warehouse.go +++ b/experimental/aitools/cmd/get_default_warehouse.go @@ -43,7 +43,7 @@ Returns warehouse ID of the default warehouse. Use --output json to get the full sess.Set(middlewares.DatabricksClientKey, w) ctx = session.WithSession(ctx, sess) - warehouse, err := middlewares.GetWarehouseEndpoint(ctx) + warehouse, err := middlewares.GetWarehouseEndpoint(ctx, false) if err != nil { return err } diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go index d1075dfd9f..23fc3c5f2d 100644 --- a/experimental/aitools/cmd/query.go +++ b/experimental/aitools/cmd/query.go @@ -41,7 +41,7 @@ Output includes the query results as JSON and row count.`, sess.Set(middlewares.DatabricksClientKey, w) ctx = session.WithSession(ctx, sess) - warehouseID, err := middlewares.GetWarehouseID(ctx) + warehouseID, err := middlewares.GetWarehouseID(ctx, true) if err != nil { return err } diff --git a/experimental/aitools/lib/middlewares/warehouse.go b/experimental/aitools/lib/middlewares/warehouse.go index dd1ab90db6..5ee3cd1387 100644 --- a/experimental/aitools/lib/middlewares/warehouse.go +++ b/experimental/aitools/lib/middlewares/warehouse.go @@ -40,7 +40,9 @@ func loadWarehouseInBackground(ctx context.Context) { sess.Set("warehouse_endpoint", warehouse) } -func GetWarehouseEndpoint(ctx context.Context) (*sql.EndpointInfo, error) { +// GetWarehouseEndpoint returns the resolved warehouse endpoint. +// If autoStart is true and the warehouse is stopped, it will be started automatically. +func GetWarehouseEndpoint(ctx context.Context, autoStart bool) (*sql.EndpointInfo, error) { sess, err := session.GetSession(ctx) if err != nil { return nil, err @@ -69,42 +71,55 @@ func GetWarehouseEndpoint(ctx context.Context) (*sql.EndpointInfo, error) { sess.Set("warehouse_endpoint", warehouse) } - return warehouse.(*sql.EndpointInfo), nil + endpoint := warehouse.(*sql.EndpointInfo) + + if autoStart && (endpoint.State == sql.StateStopped || endpoint.State == sql.StateStopping) { + endpoint, err = startWarehouse(ctx, endpoint.Id) + if err != nil { + return nil, err + } + sess.Set("warehouse_endpoint", endpoint) + } + + return endpoint, nil } -func GetWarehouseID(ctx context.Context) (string, error) { - warehouse, err := GetWarehouseEndpoint(ctx) +// GetWarehouseID returns the resolved warehouse ID. +// If autoStart is true and the warehouse is stopped, it will be started automatically. +func GetWarehouseID(ctx context.Context, autoStart bool) (string, error) { + warehouse, err := GetWarehouseEndpoint(ctx, autoStart) if err != nil { return "", err } return warehouse.Id, nil } -func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) { +func startWarehouse(ctx context.Context, id string) (*sql.EndpointInfo, error) { w, err := GetDatabricksClient(ctx) if err != nil { return nil, fmt.Errorf("get databricks client: %w", err) } - - warehouse, err := resolveWarehouse(ctx, w) + wait, err := w.Warehouses.Start(ctx, sql.StartRequest{Id: id}) if err != nil { - return nil, err + return nil, fmt.Errorf("start warehouse %s: %w", id, err) } - - // Start the warehouse if it's not running. - if warehouse.State == sql.StateStopped || warehouse.State == sql.StateStopping { - wait, err := w.Warehouses.Start(ctx, sql.StartRequest{Id: warehouse.Id}) - if err != nil { - return nil, fmt.Errorf("start warehouse %s: %w", warehouse.Id, err) - } - resp, err := wait.Get() - if err != nil { - return nil, fmt.Errorf("wait for warehouse %s to start: %w", warehouse.Id, err) - } - warehouse.State = resp.State + resp, err := wait.Get() + if err != nil { + return nil, fmt.Errorf("wait for warehouse %s to start: %w", id, err) } + return &sql.EndpointInfo{ + Id: resp.Id, + Name: resp.Name, + State: resp.State, + }, nil +} - return warehouse, nil +func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) { + w, err := GetDatabricksClient(ctx) + if err != nil { + return nil, fmt.Errorf("get databricks client: %w", err) + } + return resolveWarehouse(ctx, w) } // resolveWarehouse selects a warehouse using the following priority: diff --git a/experimental/aitools/lib/providers/clitools/discover.go b/experimental/aitools/lib/providers/clitools/discover.go index 7b67b2a6fa..5c69ba16a7 100644 --- a/experimental/aitools/lib/providers/clitools/discover.go +++ b/experimental/aitools/lib/providers/clitools/discover.go @@ -16,7 +16,7 @@ import ( // Discover provides workspace context and workflow guidance. // Returns L1 (flow) always + L2 (target) for detected target types + L3 (skills) listing. func Discover(ctx context.Context, workingDirectory string) (string, error) { - warehouse, err := middlewares.GetWarehouseEndpoint(ctx) + warehouse, err := middlewares.GetWarehouseEndpoint(ctx, false) if err != nil { log.Debugf(ctx, "Failed to get default warehouse (non-fatal): %v", err) warehouse = nil