Skip to content

Commit 6fbc7af

Browse files
fjakobsclaude
andauthored
Add get-default-warehouse command to aitools (#4242)
## Why This is something we need for experimenting with skills. There is currently no command that finds the default waregouse. ## Summary - Add `databricks experimental aitools tools get-default-warehouse` command - Returns warehouse ID in text mode (default) for easy scripting - Returns full warehouse info (id, name, state) in JSON mode with `--output json` - Uses existing `GetWarehouseID` function from middlewares package ## Examples ```bash # Text output (default) databricks experimental aitools tools get-default-warehouse # Output: abc123def456... # JSON output databricks experimental aitools tools get-default-warehouse --output json # Output: {"id":"abc123def456...","name":"My Warehouse","state":"RUNNING"} ``` 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Sonnet 4.5 <[email protected]>
1 parent c5b7c02 commit 6fbc7af

File tree

2 files changed

+63
-0
lines changed

2 files changed

+63
-0
lines changed
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package mcp
2+
3+
import (
4+
"github.com/databricks/cli/cmd/root"
5+
"github.com/databricks/cli/experimental/aitools/lib/middlewares"
6+
"github.com/databricks/cli/experimental/aitools/lib/session"
7+
"github.com/databricks/cli/libs/cmdctx"
8+
"github.com/databricks/cli/libs/cmdio"
9+
"github.com/databricks/databricks-sdk-go/service/sql"
10+
"github.com/spf13/cobra"
11+
)
12+
13+
type warehouseInfo struct {
14+
Id string `json:"id"`
15+
Name string `json:"name"`
16+
State sql.State `json:"state"`
17+
}
18+
19+
func newGetDefaultWarehouseCmd() *cobra.Command {
20+
cmd := &cobra.Command{
21+
Use: "get-default-warehouse",
22+
Short: "Get the default warehouse ID",
23+
Long: `Get the default warehouse ID for the current workspace.
24+
25+
The command auto-detects an available warehouse unless DATABRICKS_WAREHOUSE_ID is set.
26+
27+
Returns warehouse ID of the default warehouse. Use --output json to get the full warehouse info including name and state.`,
28+
Example: ` # Get warehouse ID in text format (default)
29+
databricks experimental aitools tools get-default-warehouse
30+
# Output: abc123def456...
31+
32+
# Get full warehouse info including name and state in JSON format
33+
databricks experimental aitools tools get-default-warehouse --output json
34+
# Output: {"id":"abc123def456...","name":"My Warehouse","state":"RUNNING"}`,
35+
Args: cobra.NoArgs,
36+
PreRunE: root.MustWorkspaceClient,
37+
RunE: func(cmd *cobra.Command, args []string) error {
38+
ctx := cmd.Context()
39+
w := cmdctx.WorkspaceClient(ctx)
40+
41+
// set up session with client for middleware compatibility
42+
sess := session.NewSession()
43+
sess.Set(middlewares.DatabricksClientKey, w)
44+
ctx = session.WithSession(ctx, sess)
45+
46+
warehouse, err := middlewares.GetWarehouseEndpoint(ctx)
47+
if err != nil {
48+
return err
49+
}
50+
51+
info := warehouseInfo{
52+
Id: warehouse.Id,
53+
Name: warehouse.Name,
54+
State: warehouse.State,
55+
}
56+
57+
return cmdio.RenderWithTemplate(ctx, info, "", "{{.Id}}\n")
58+
},
59+
}
60+
61+
return cmd
62+
}

experimental/aitools/cmd/tools.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ func newToolsCmd() *cobra.Command {
1717
cmd.AddCommand(init_template.NewInitTemplateCommand())
1818
cmd.AddCommand(newValidateCmd())
1919
cmd.AddCommand(newDeployCmd())
20+
cmd.AddCommand(newGetDefaultWarehouseCmd())
2021

2122
return cmd
2223
}

0 commit comments

Comments
 (0)