diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2f02fb16ea..2b4f85a83a 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -7,5 +7,4 @@ /cmd/workspace/apps/ @databricks/eng-app-devex /libs/apps/ @databricks/eng-app-devex /acceptance/apps/ @databricks/eng-app-devex -/experimental/aitools/ @databricks/eng-app-devex @lennartkats-db /experimental/apps-mcp/ @databricks/eng-app-devex @lennartkats-db diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 5533af9065..defcebced0 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -145,37 +145,6 @@ jobs: - name: Analyze slow tests run: make slowest - test-exp-aitools: - needs: - - cleanups - - testmask - - # Only run if the target is in the list of targets from testmask - if: ${{ contains(fromJSON(needs.testmask.outputs.targets), 'test-exp-aitools') }} - name: "make test-exp-aitools" - runs-on: ${{ matrix.os }} - - strategy: - fail-fast: false - matrix: - os: - - macos-latest - - ubuntu-latest - - windows-latest - - steps: - - name: Checkout repository and submodules - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - - - name: Setup build environment - uses: ./.github/actions/setup-build-environment - with: - cache-key: test-exp-aitools - - - name: Run tests - run: | - make test-exp-aitools - test-exp-apps-mcp: needs: - cleanups @@ -282,7 +251,6 @@ jobs: test-result: needs: - test - - test-exp-aitools - test-exp-apps-mcp - test-exp-ssh - test-pipelines diff --git a/Makefile b/Makefile index 031ad94c3b..92bbf49332 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ default: checks fmt lint # Default packages to test (all) -TEST_PACKAGES = ./acceptance/internal ./libs/... ./internal/... ./cmd/... ./bundle/... ./experimental/aitools/... ./experimental/ssh/... . +TEST_PACKAGES = ./acceptance/internal ./libs/... ./internal/... ./cmd/... ./bundle/... ./experimental/ssh/... . # Default acceptance test filter (all) ACCEPTANCE_TEST_FILTER = "" @@ -173,9 +173,6 @@ generate: .PHONY: lint lintfull tidy lintcheck fmt fmtfull test test-unit test-acc test-slow test-slow-unit test-slow-acc cover showcover build snapshot snapshot-release schema integration integration-short acc-cover acc-showcover docs ws wsfix links checks test-update test-update-templates generate-out-test-toml test-update-aws test-update-all generate-validation -test-exp-aitools: - make test TEST_PACKAGES="./experimental/aitools/..." ACCEPTANCE_TEST_FILTER="TestAccept/idontexistyet/aitools" - test-exp-apps-mcp: make test TEST_PACKAGES="./experimental/apps-mcp/..." ACCEPTANCE_TEST_FILTER="TestAccept/idontexistyet/apps-mcp" diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index 79712b5679..3bdd38170c 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -1,7 +1,6 @@ package experimental import ( - "github.com/databricks/cli/experimental/aitools" mcp "github.com/databricks/cli/experimental/apps-mcp/cmd" "github.com/spf13/cobra" ) @@ -21,7 +20,6 @@ These commands provide early access to new features that are still under development. They may change or be removed in future versions without notice.`, } - cmd.AddCommand(aitools.New()) cmd.AddCommand(mcp.NewMcpCmd()) return cmd diff --git a/experimental/aitools/DESIGN.md b/experimental/aitools/DESIGN.md deleted file mode 100644 index 7cc141858a..0000000000 --- a/experimental/aitools/DESIGN.md +++ /dev/null @@ -1,227 +0,0 @@ -Add a 'databricks experimental aitools' command that behaves as follows: - -* databricks experimental aitools --help: shows help -* databricks experimental aitools: shows help (like other command groups) -* databricks experimental aitools install: installs the server in coding agents -* databricks experimental aitools server: starts the mcp server (subcommand) -* databricks experimental aitools uninstall: uninstalls the server from coding agents (subcommand - not implemented; errors out and tells the user to ask their local coding agent to uninstall the Databricks CLI MCP server) -* databricks experimental aitools tool --json : runs a specific MCP tool for acceptance testing (hidden subcommand) - -non-functional requirements: -- any errors that these commands give should be friendly, concise, actionable. -- this code needs to be modular (e.g cursor installation/detection is one module) and needs to have unit tests -- write code docs and documentation in a very concise and minimal way, and keep maintainers in mind; look at other - modules for inspiration -- take AGENTS.md into account when building this -- MANDATORY: never invoke the databricks cli directly, instead use the invoke_databricks_cli tool! -- all AI-facing prompts and guidance text must be stored in .tmpl files in the tools/prompts/ directory (not hardcoded in Go files) -- Resource-specific code is modularized in cmd/mcp/tools/resources/ directory. Each resource type (app, job, pipeline, dashboard) has its own file (apps.go, jobs.go, etc.) that implements: - - Add* function: handles adding the resource to a project - - AnalyzeGuidance* function: returns resource-specific guidance for the AI agent - This structure makes it easy for teams (e.g., the apps team) to customize how the agent works with their specific resource type by editing the corresponding resource file. -- MANDATORY: you need to experiment locally with the claude cli and cursor cli to make sure this actually works as expected. - The testing approach should be: - - Example test command: - ```bash - rm -rf /tmp/blank; mkdir -p /tmp/blank; cd /tmp/blank; - claude --allow-dangerously-skip-permissions "Create a new Databricks app that shows a dashboard with taxi trip fares per city, then preview it and open it in my browser. If the databricks-aitools MCP fails, stop immediately and ask for my guidance." - ``` - - You should test multiple scenarios: - - Creating a project with a simple job that lists taxis and running it - - Creating an app with a dashboard - - Adding resources to an existing project - - The key is to use the Claude CLI to issue prompts as a user would, NOT to directly call MCP tools yourself - - If the MCP server has issues, Claude Code should surface clear error messages - - This is the most important part of the work; i expect you to deeply experiment with the new mcp server; add it, try it, remove it, add it again, use it to build things that can run. -- MANDATORY: at the very end, compare what you built again to the instructions in this doc; go over each point, does it work, is it complete? - - -To illustrate how the install command works: - -``` -$ databricks experimental aitools install - - ▄▄▄▄▄▄▄▄ Databricks CLI - ██▌ ▐██ MCP Server - ▀▀▀▀▀▀▀▀ - -╔════════════════════════════════════════════════════════════════╗ -║ ⚠️ EXPERIMENTAL: This command may change in future versions ║ -╚════════════════════════════════════════════════════════════════╝ - -Which coding agents would you like to install the MCP server for? - -Install for Claude Code? [Use arrows to move] -> yes - no - -Install for Cursor? [Use arrows to move] -> yes - no - -Show manual installation instructions for other agents? [Use arrows to move] -> no - yes - -Installing MCP server for Claude Code... -✓ Installed for Claude Code - -Installing MCP server for Cursor... -✓ Installed for Cursor - -You can now use your coding agent to interact with Databricks. -Try asking: 'Create a new Databricks project with a job or an app' -``` - -Implementation notes: -- Uses cmdio.AskSelect (arrow keys, all default to "yes" for easy Enter×3) -- Always prompts for both Claude Code and Cursor (no detection filtering) -- Gracefully skips missing agents with yellow warning instead of erroring -- Line replacement: "Installing..." → "✓ Installed" (or "⊘ Skipped") using \r - -Now databricks experimental aitools server should actually start an MCP server that we actually use to describe -the system as a whole a bit (btw each tool should be defined in a separate .go file, right!): -- when starting up it should do a the 'roots/list' to the agent. - - if that doesn't work or if there is more than one root => error out - - look at the root path. if there is already a databricks.yml file, that means the user already initialized a project; keep track of that. - -- for the tools below, there is a common initialization step. you need to check if you're - properly authenticated to the workspace. you can do so by using the invoke_databricks_cli tool - to run 'jobs get ' (pick any random id like 123456789). - if you get an authentication error, the tools need to return an error saying that they first need - to log in to databricks. to do so they need to use the invoke_databricks_cli tool to run: - 'auth login --profile DEFAULT --host '. - the AI needs to ask the user for this url, it cannot guess it. - once logged in, the tools will work! the AI should also point to https://docs.databricks.com/getting-started/free-edition as a page where users can setup their own fully free account for experimentation. - -- the "invoke_databricks_cli" tool: - - description: run any databricks CLI command. this is a passthrough to the databricks CLI. - use this tool whenever you need to run databricks CLI commands like 'databricks bundle deploy', - 'databricks bundle validate', 'databricks bundle run', 'databricks auth login', etc. - the reason this tool exists (instead of invoking the databricks CLI directly) is to make it - easier for users to allow-list commands compared to allow-listing shell commands. - - parameter: command - the full databricks CLI command to run, e.g. "bundle deploy" or "bundle validate" - (note: do not include the "databricks" prefix in the command parameter) - - parameter: working_directory - optional. the directory to run the command in. defaults to the current directory. - - output: the stdout and stderr from the command, plus the exit code - - implementation guidance: this should just invoke the databricks CLI and return the output. - make sure to properly handle the working directory if provided. - - further implementation guidance: i want an acceptance test for this command. it should just call the 'help' command. - -- the "init_project" tool: - - description: initializes a new databricks project structure. Use this to create a new project. After initialization, use add_project_resource to add resources such as apps, jobs, dashboards, pipelines, etc. - - parameter: project_name - a name for this project in snake_case; ask the user about this if it's not clear from the context - - parameter: project_path - a fully qualified path for the directory to create the new project in. Usually this should be in the current directory! But if it already has a lot of other things then it should be a subdirectory. - - action to perform when this runs: use the invoke_databricks_cli tool to run - 'bundle init default-minimal --config-file /tmp/...' where you set the 'project_name' and other - parameters. use personal schemas and the default catalog. - note that default-minimal creates a subdirectory called 'project_name'! this is not needed. just move everything - (recursively) in that subdirectory to the target directory from project_path. - after initialization, creates a CLAUDE.md file (if the calling MCP client is Claude Code) or AGENTS.md file (otherwise) - with project-specific agent instructions. The file includes: - - Installation instructions for the Databricks CLI MCP server (if not yet installed) - - Guidance to use the mcp__databricks-aitools__analyze_project tool when opening the project - The client is detected at runtime from the MCP initialize request's clientInfo field. - - guidance on how to implement this: do some trial and error to make the init command work. - do a non-forward merge of origin/add-default-minimal to get the minimal template! - - output: returns a success message with a WARNING that this is an EMPTY project with NO resources, and that add_project_resource MUST be called if the user asked for a specific resource. followed by the same guidance as the analyze_project tool (calls analyze_project internally) - - further implementation guidance: i want an acceptance test for this command. it should lead to a project - that can pass a 'bundle validate' command! - -- the "analyze_project" tool: - - description: REQUIRED FIRST STEP: If databricks.yml exists in the directory, you MUST call this tool before using Read, Glob, or any other tools. Databricks projects require specialized commands that differ from standard Python/Node.js workflows - attempting standard approaches will fail. This tool is fast and provides the correct commands for preview/deploy/run operations. - - parameter: project_path - a fully qualified path of the project to operate on. - - output: - - summary: contents of the 'bundle summary' command run in this dir using the invoke_databricks_cli tool. - - - guidance: - - "Below is guidance for how to work with this project. - - IMPORTANT: you want to give the user some idea of how to get started; you can suggest - prompts such as "Create an app that shows a chart with taxi fares by city" - or "Create a job that summarizes all taxi data using a notebook" - - IMPORTANT: Most interactions are done with the Databricks CLI. YOU (the AI) must use the invoke_databricks_cli tool to run commands - never suggest the user runs CLI commands directly! - - IMPORTANT: to add new resources to a project, use the 'add_project_resource' mcp tool. - - MANDATORY: always deploy with invoke_databricks_cli 'bundle deploy', never with 'apps deploy' - - Note that Databricks resources are defined in resources/*.yml files. See https://docs.databricks.com/dev-tools/bundles/settings for a reference! - - - Common guidance about getting started and using the CLI (should draw inspiration from the original default_python_template_readme.md file, extracting common instructions that are not app-specific) - - /README.md (if it exists). This provides project-specific guidance that complements the common guidance> - -- the "add_project_resource" tool: - - description: extend the current project with a new app, job, pipeline, or dashboard - - parameter: type - app, job, pipeline, or dashboard - - parameter: name - the name of the new resource (for example: new_app); should not exist yet in resources/ - - parameter: template - optional. only fill this in when asked. - - implementation guidance: - - (i have some idea how to implement apps, as described below, but for now just error say its not implemented) - for apps, there are templates in https://github.com/databricks/app-templates. - - if no template was given, error out and tell the AI: either pick a template from the list of templates, or let the user pick. if the user didn't pick a template but did describe an app then just - use nodejs-fastapi-hello-world-app as a starting point. - - if a template _was_ given then you should create a shallow git clone of https://github.com/databricks/app-templates in /tmp and then copy one of the template dirs (e.g. e2e-chatbot-app-next) to a folder with that name (e.g. e2e-chatbot-app-next). you should also create an associated resources/*.yml (e2e-chatbot-app-next.yml) see https://github.com/databricks/bundle-examples/blob/main/knowledge_base/databricks_app/resources/hello_world.job.yml for a starting point. - - for jobs, the template parameter needs to be sql or python. error out if not specified; the ai needs to - ask what language the user wants if this was not clear from the context. - if a template is specified then do a shallow clone of https://github.com/databricks/bundle-examples, - and take default_python or default_sql as a starting point depending on the language. - you need to copy resources/*.job.yml but rename them to resources/.job.yml - for python, you need to copy src/default_python (but rename to src/) and src/tests - for sql, you need to copy src/*.sql (dont overwrite anything) - - for pipelines, most of the guidance is the same (the implementation could be shared?): - the template parameter needs to be sql or python. error out if not specified; the ai needs to - ask what language the user wants if this was not clear from the context. - if a template is specified then do a shallow clone of https://github.com/databricks/bundle-examples, - and take lakeflow_pipelines_python or lakeflow_pipelines_sql as a starting point depending on the language. - you need to copy resources/*.pipeline.yml but rename them to resources/.pipeline.yml - copy src/lakeflow_pipeline_*/** but rename to src//** - - for dashboards, do a shallow clone of https://github.com/databricks/bundle-examples. - use knowledge_base/dashboard_nyc_taxi as a starting point. - you need to copy resources/*.dashboard.yml but rename them to resources/.dashboard.yml - copy src/*.lvdash.json but rename to src/.lvdash.json - - note that all of the above (apps, jobs, pipelines, dashboards) should include a note that says "FIXME: this should rely on the databricks bundle generate command" - - output: if any of the resource types from above (e.g. a python job) were instantiated, - the output needs to respond with guidance yet. it should say that the MCP only - created a starting point and that the AI needs to iterate over it: - 1. use the analyze_project tool to learn about the current project structure and how to use it - 2. validate that the extensions are correct using the invoke_databricks_cli tool to run 'bundle validate'. for apps, also check that any warehouse references in the resource/*.yml file are valid. - 3. should fix any errors, and eventually should deploy to dev using the invoke_databricks_cli tool to run 'bundle deploy --target dev' - for a pipeline, it can also use the invoke_databricks_cli tool to run 'bundle run --validate-only' to - validate that the pipeline is correct. - - further implementation guidance: i want acceptance tests for each of these project types (app, dashboard, job, pipeline); - this means they should be exposed as a hidden command like 'databricks experimental aitools tool add_project_resource --json '. having these tests will be instrumental for iterating on them; the initing should not fail! note that the tool subcommand should just assume that the cwd is the current project dir. - -- the "explore" tool: - - description: CALL THIS FIRST when user mentions a workspace by name or asks about workspace resources. Shows available workspaces/profiles, default warehouse, and provides guidance on exploring jobs, clusters, catalogs, and other Databricks resources. Use this to discover what's available before running CLI commands. - - no parameters needed - - implementation: - - Determines a default SQL warehouse for queries using GetDefaultWarehouse(): - 1. Lists all warehouses using 'databricks warehouses list --output json' - 2. Prefers RUNNING warehouses (pick first one found) - 3. If none running, picks first STOPPED warehouse (warehouses auto-start when queried) - 4. If no warehouses available, returns error directing user to create one - - Shows workspace/profile information: - 1. Reads available profiles from ~/.databrickscfg using libs/databrickscfg/profile package - 2. Shows current profile (from DATABRICKS_CONFIG_PROFILE env var or DEFAULT) - 3. Lists all available workspaces with their host URLs and cloud providers - 4. Provides guidance on using --profile flag to switch workspaces: - - Example: invoke_databricks_cli '--profile prod catalogs list' - 5. Only shows profile list if multiple profiles exist (saves tokens for single-profile setups) - - Returns concise guidance text that explains: - 1. Current workspace profile and host - 2. Available workspace profiles (if multiple exist) - 3. The warehouse ID that can be used for queries - 4. How to execute SQL queries using Statement Execution API: - - invoke_databricks_cli 'api post /api/2.0/sql/statements --json {"warehouse_id":"...","statement":"SELECT ...","wait_timeout":"30s"}' - - Mentions using the warehouse ID shown above - 5. How to explore workspace resources: - - Jobs: invoke_databricks_cli 'jobs list', 'jobs get ' - - Clusters: invoke_databricks_cli 'clusters list', 'clusters get ' - - Unity Catalog: invoke_databricks_cli 'catalogs list', 'schemas list', 'tables list', 'tables get' - - Workspace files: invoke_databricks_cli 'workspace list ' - 6. Reminder to use --profile flag for non-default workspaces - - Key design: Single concise endpoint that provides guidance, not many separate tools - - output: Guidance text with workspace/profile info, warehouse info, and commands for exploring jobs, clusters, data, and other resources - - implementation: Single explore.go file with GetDefaultWarehouse, getCurrentProfile, and getAvailableProfiles helpers - - key use case: When user asks about a specific workspace (e.g., "what jobs do I have in my dogfood workspace"), agent should call this FIRST to see available workspaces and get the correct profile name diff --git a/experimental/aitools/agents/claude.go b/experimental/aitools/agents/claude.go deleted file mode 100644 index 7c0b335b81..0000000000 --- a/experimental/aitools/agents/claude.go +++ /dev/null @@ -1,44 +0,0 @@ -package agents - -import ( - "errors" - "fmt" - "os/exec" - - "github.com/databricks/cli/experimental/aitools/tools" -) - -// DetectClaude checks if Claude Code CLI is installed and available on PATH. -func DetectClaude() bool { - _, err := exec.LookPath("claude") - return err == nil -} - -// InstallClaude installs the Databricks MCP server in Claude Code. -func InstallClaude() error { - if !DetectClaude() { - return errors.New("claude Code CLI is not installed or not on PATH\n\nPlease install Claude Code and ensure 'claude' is available on your system PATH.\nFor installation instructions, visit: https://docs.anthropic.com/en/docs/claude-code") - } - - databricksPath, err := tools.GetDatabricksPath() - if err != nil { - return err - } - - removeCmd := exec.Command("claude", "mcp", "remove", "--scope", "user", "databricks-aitools") - _ = removeCmd.Run() - - cmd := exec.Command("claude", "mcp", "add", - "--scope", "user", - "--transport", "stdio", - "databricks-aitools", - "--", - databricksPath, "experimental", "aitools", "server") - - output, err := cmd.CombinedOutput() - if err != nil { - return fmt.Errorf("failed to install MCP server in Claude Code: %w\nOutput: %s", err, string(output)) - } - - return nil -} diff --git a/experimental/aitools/agents/cursor.go b/experimental/aitools/agents/cursor.go deleted file mode 100644 index 62547ef3a5..0000000000 --- a/experimental/aitools/agents/cursor.go +++ /dev/null @@ -1,96 +0,0 @@ -package agents - -import ( - "encoding/json" - "fmt" - "os" - "path/filepath" - "runtime" -) - -type cursorConfig struct { - McpServers map[string]mcpServer `json:"mcpServers"` -} - -type mcpServer struct { - Command string `json:"command"` - Args []string `json:"args,omitempty"` - Env map[string]string `json:"env,omitempty"` -} - -func getCursorConfigPath() (string, error) { - if runtime.GOOS == "windows" { - userProfile := os.Getenv("USERPROFILE") - if userProfile == "" { - return "", os.ErrNotExist - } - return filepath.Join(userProfile, ".cursor", "mcp.json"), nil - } - - home, err := os.UserHomeDir() - if err != nil { - return "", err - } - return filepath.Join(home, ".cursor", "mcp.json"), nil -} - -// DetectCursor checks if Cursor is installed by looking for its config directory. -func DetectCursor() bool { - configPath, err := getCursorConfigPath() - if err != nil { - return false - } - // Check if the .cursor directory exists (not just the mcp.json file) - cursorDir := filepath.Dir(configPath) - _, err = os.Stat(cursorDir) - return err == nil -} - -// InstallCursor installs the Databricks MCP server in Cursor. -func InstallCursor() error { - configPath, err := getCursorConfigPath() - if err != nil { - return fmt.Errorf("failed to determine Cursor config path: %w", err) - } - - // Check if .cursor directory exists (not the file, we'll create that if needed) - cursorDir := filepath.Dir(configPath) - if _, err := os.Stat(cursorDir); err != nil { - return fmt.Errorf("cursor directory not found at: %s\n\nPlease install Cursor from: https://cursor.sh", cursorDir) - } - - // Read existing config - var config cursorConfig - data, err := os.ReadFile(configPath) - if err != nil { - // If file doesn't exist or can't be read, start with empty config - config = cursorConfig{ - McpServers: make(map[string]mcpServer), - } - } else { - if err := json.Unmarshal(data, &config); err != nil { - return fmt.Errorf("failed to parse Cursor config: %w", err) - } - if config.McpServers == nil { - config.McpServers = make(map[string]mcpServer) - } - } - - // Add or update the Databricks MCP server entry - config.McpServers["databricks"] = mcpServer{ - Command: "databricks", - Args: []string{"experimental", "aitools", "server"}, - } - - // Write back to file with pretty printing - updatedData, err := json.MarshalIndent(config, "", " ") - if err != nil { - return fmt.Errorf("failed to marshal Cursor config: %w", err) - } - - if err := os.WriteFile(configPath, updatedData, 0o644); err != nil { - return fmt.Errorf("failed to write Cursor config: %w", err) - } - - return nil -} diff --git a/experimental/aitools/agents/custom.go b/experimental/aitools/agents/custom.go deleted file mode 100644 index 27eee1d55e..0000000000 --- a/experimental/aitools/agents/custom.go +++ /dev/null @@ -1,36 +0,0 @@ -package agents - -import ( - "context" - - "github.com/databricks/cli/libs/cmdio" -) - -// ShowCustomInstructions displays instructions for manually installing the MCP server. -func ShowCustomInstructions(ctx context.Context) error { - instructions := ` -To install the Databricks CLI MCP server in your coding agent: - -1. Add a new MCP server to your coding agent's configuration -2. Set the command to: databricks experimental aitools server -3. No environment variables or additional configuration needed - -Example MCP server configuration: -{ - "mcpServers": { - "databricks": { - "command": "databricks", - "args": ["experimental", "aitools", "server"] - } - } -} -` - cmdio.LogString(ctx, instructions) - - _, err := cmdio.Ask(ctx, "Press Enter to continue", "") - if err != nil { - return err - } - cmdio.LogString(ctx, "") - return nil -} diff --git a/experimental/aitools/aitools.go b/experimental/aitools/aitools.go deleted file mode 100644 index d61183e5d3..0000000000 --- a/experimental/aitools/aitools.go +++ /dev/null @@ -1,34 +0,0 @@ -package aitools - -import ( - "github.com/spf13/cobra" -) - -func New() *cobra.Command { - cmd := &cobra.Command{ - Use: "aitools", - Short: "Manage AI agent skills and MCP server for coding agents", - Hidden: true, - Long: `Manage Databricks AI agent skills and MCP server. - -This provides AI agents like Claude Code and Cursor with capabilities to interact -with Databricks, create projects, deploy bundles, run jobs, etc. - -╔════════════════════════════════════════════════════════════════╗ -║ ⚠️ EXPERIMENTAL: This command may change in future versions ║ -╚════════════════════════════════════════════════════════════════╝ - -Common workflows: - databricks experimental aitools install # Install in Claude Code or Cursor - databricks experimental aitools server # Start server (used by agents) - -Online documentation: https://docs.databricks.com/dev-tools/cli/aitools.html`, - } - - cmd.AddCommand(newInstallCmd()) - cmd.AddCommand(newServerCmd()) - cmd.AddCommand(newUninstallCmd()) - cmd.AddCommand(newToolCmd()) - - return cmd -} diff --git a/experimental/aitools/auth/check.go b/experimental/aitools/auth/check.go deleted file mode 100644 index 609fd06542..0000000000 --- a/experimental/aitools/auth/check.go +++ /dev/null @@ -1,67 +0,0 @@ -package auth - -import ( - "context" - "errors" - "net/http" - "os" - "sync" - - "github.com/databricks/cli/experimental/aitools/tools/prompts" - "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/service/jobs" -) - -var ( - authCheckOnce sync.Once - authCheckResult error -) - -// CheckAuthentication checks if the user is authenticated to a Databricks workspace. -// It caches the result so the check only runs once per process. -func CheckAuthentication(ctx context.Context) error { - authCheckOnce.Do(func() { - authCheckResult = checkAuth(ctx) - }) - return authCheckResult -} - -func checkAuth(ctx context.Context) error { - if os.Getenv("DATABRICKS_MCP_SKIP_AUTH_CHECK") == "1" { - return nil - } - - w, err := databricks.NewWorkspaceClient() - if err != nil { - return wrapAuthError(err) - } - - // Use Jobs API for auth check (fast). Expected: 404 (authenticated), 401/403 (not authenticated). - _, err = w.Jobs.Get(ctx, jobs.GetJobRequest{JobId: 999999999}) - if err == nil { - return nil - } - - var apiErr *apierr.APIError - if errors.As(err, &apiErr) { - switch apiErr.StatusCode { - case http.StatusNotFound: - return nil - case http.StatusUnauthorized, http.StatusForbidden: - return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil)) - default: - return nil - } - } - - return wrapAuthError(err) -} - -func wrapAuthError(err error) error { - if errors.Is(err, config.ErrCannotConfigureDefault) { - return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil)) - } - return err -} diff --git a/experimental/aitools/install.go b/experimental/aitools/install.go deleted file mode 100644 index 7feabe3c24..0000000000 --- a/experimental/aitools/install.go +++ /dev/null @@ -1,96 +0,0 @@ -package aitools - -import ( - "context" - "fmt" - "os" - "time" - - "github.com/databricks/cli/experimental/aitools/agents" - "github.com/databricks/cli/libs/cmdio" - "github.com/fatih/color" - "github.com/spf13/cobra" -) - -func newInstallCmd() *cobra.Command { - cmd := &cobra.Command{ - Use: "install", - Short: "Install the MCP server in coding agents", - Long: `Install the Databricks CLI MCP server in coding agents like Claude Code and Cursor.`, - RunE: func(cmd *cobra.Command, args []string) error { - return runInstall(cmd.Context()) - }, - } - - return cmd -} - -func runInstall(ctx context.Context) error { - cmdio.LogString(ctx, "") - green := color.New(color.FgGreen).SprintFunc() - cmdio.LogString(ctx, " "+green("[")+"████████"+green("]")+" Databricks Experimental AI agent skills & MCP") - cmdio.LogString(ctx, " "+green("[")+"██▌ ▐██"+green("]")) - cmdio.LogString(ctx, " "+green("[")+"████████"+green("]")+" AI-powered Databricks development and exploration") - cmdio.LogString(ctx, "") - - yellow := color.New(color.FgYellow).SprintFunc() - cmdio.LogString(ctx, yellow("╔════════════════════════════════════════════════════════════════╗")) - cmdio.LogString(ctx, yellow("║ ⚠️ EXPERIMENTAL: This command may change in future versions ║")) - cmdio.LogString(ctx, yellow("╚════════════════════════════════════════════════════════════════╝")) - cmdio.LogString(ctx, "") - - cmdio.LogString(ctx, "Which coding agents would you like to install the MCP server for?") - cmdio.LogString(ctx, "") - - anySuccess := false - - ans, err := cmdio.AskSelect(ctx, "Install for Claude Code?", []string{"yes", "no"}) - if err != nil { - return err - } - if ans == "yes" { - fmt.Fprint(os.Stderr, "Installing MCP server for Claude Code...") - if err := agents.InstallClaude(); err != nil { - fmt.Fprint(os.Stderr, "\r"+color.YellowString("⊘ Skipped Claude Code: "+err.Error())+"\n") - } else { - fmt.Fprint(os.Stderr, "\r"+color.GreenString("✓ Installed for Claude Code")+" \n") - anySuccess = true - } - cmdio.LogString(ctx, "") - } - - ans, err = cmdio.AskSelect(ctx, "Install for Cursor?", []string{"yes", "no"}) - if err != nil { - return err - } - if ans == "yes" { - fmt.Fprint(os.Stderr, "Installing MCP server for Cursor...") - if err := agents.InstallCursor(); err != nil { - fmt.Fprint(os.Stderr, "\r"+color.YellowString("⊘ Skipped Cursor: "+err.Error())+"\n") - } else { - // Brief delay so users see the "Installing..." message before it's replaced - time.Sleep(1 * time.Second) - fmt.Fprint(os.Stderr, "\r"+color.GreenString("✓ Installed for Cursor")+" \n") - anySuccess = true - } - cmdio.LogString(ctx, "") - } - - ans, err = cmdio.AskSelect(ctx, "Show manual installation instructions for other agents?", []string{"yes", "no"}) - if err != nil { - return err - } - if ans == "yes" { - if err := agents.ShowCustomInstructions(ctx); err != nil { - return err - } - } - - if anySuccess { - cmdio.LogString(ctx, "") - cmdio.LogString(ctx, "You can now use your coding agent to interact with Databricks.") - cmdio.LogString(ctx, "Try asking: 'Create a new Databricks project with a job or an app'") - } - - return nil -} diff --git a/experimental/aitools/server.go b/experimental/aitools/server.go deleted file mode 100644 index 7d9268b0a3..0000000000 --- a/experimental/aitools/server.go +++ /dev/null @@ -1,254 +0,0 @@ -package aitools - -import ( - "bufio" - "context" - "encoding/json" - "fmt" - "io" - "os" - - "github.com/databricks/cli/experimental/aitools/tools" - "github.com/databricks/cli/libs/log" - "github.com/spf13/cobra" -) - -func newServerCmd() *cobra.Command { - cmd := &cobra.Command{ - Use: "server", - Short: "Start the MCP server (used by coding agents)", - Long: `Start the Databricks CLI MCP server. This command is typically invoked by coding agents.`, - RunE: func(cmd *cobra.Command, args []string) error { - return runServer(cmd.Context()) - }, - } - - return cmd -} - -type jsonrpcRequest struct { - JSONRPC string `json:"jsonrpc"` - ID any `json:"id,omitempty"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` -} - -type jsonrpcResponse struct { - JSONRPC string `json:"jsonrpc"` - ID any `json:"id,omitempty"` - Result any `json:"result,omitempty"` - Error *rpcError `json:"error,omitempty"` -} - -type rpcError struct { - Code int `json:"code"` - Message string `json:"message"` - Data any `json:"data,omitempty"` -} - -// JSON-RPC 2.0 error codes. -const ( - jsonRPCParseError = -32700 - jsonRPCInvalidRequest = -32600 - jsonRPCMethodNotFound = -32601 - jsonRPCInvalidParams = -32602 - jsonRPCInternalError = -32603 -) - -type mcpServer struct { - ctx context.Context - in io.Reader - out io.Writer - toolsMap map[string]tools.ToolHandler - clientName string -} - -// getAllTools returns all tools (definitions + handlers) for the MCP server. -func getAllTools() []tools.Tool { - return []tools.Tool{ - tools.InvokeDatabricksCLITool, - tools.InitProjectTool, - tools.AnalyzeProjectTool, - tools.AddProjectResourceTool, - tools.ExploreTool, - } -} - -// NewMCPServer creates a new MCP server instance. -func NewMCPServer(ctx context.Context) *mcpServer { - allTools := getAllTools() - toolsMap := make(map[string]tools.ToolHandler, len(allTools)) - for _, tool := range allTools { - toolsMap[tool.Definition.Name] = tool.Handler - } - - return &mcpServer{ - ctx: ctx, - in: os.Stdin, - out: os.Stdout, - toolsMap: toolsMap, - } -} - -// Start starts the MCP server and processes requests. -// Note: No logging in server mode as it interferes with JSON-RPC over stdout/stdin. -func (s *mcpServer) Start() error { - scanner := bufio.NewScanner(s.in) - for scanner.Scan() { - line := scanner.Text() - if line == "" { - continue - } - - var req jsonrpcRequest - if err := json.Unmarshal([]byte(line), &req); err != nil { - s.sendError(nil, jsonRPCParseError, "Parse error", nil) - continue - } - - s.handleRequest(&req) - } - - if err := scanner.Err(); err != nil { - return fmt.Errorf("error reading from stdin: %w", err) - } - - return nil -} - -// handleRequest processes an incoming JSON-RPC request. -func (s *mcpServer) handleRequest(req *jsonrpcRequest) { - switch req.Method { - case "initialize": - s.handleInitialize(req) - case "tools/list": - s.handleToolsList(req) - case "tools/call": - s.handleToolsCall(req) - default: - s.sendError(req.ID, jsonRPCMethodNotFound, "Method not found", nil) - } -} - -// handleInitialize handles the initialize request. -func (s *mcpServer) handleInitialize(req *jsonrpcRequest) { - // Parse clientInfo from the request - var params struct { - ClientInfo struct { - Name string `json:"name"` - Version string `json:"version"` - } `json:"clientInfo"` - } - if req.Params != nil { - _ = json.Unmarshal(req.Params, ¶ms) - s.clientName = params.ClientInfo.Name - } - - result := map[string]any{ - "protocolVersion": "2024-11-05", - "serverInfo": map[string]string{ - "name": "databricks-aitools", - "version": "1.0.0", - }, - "capabilities": map[string]any{ - "tools": map[string]bool{}, - }, - } - - s.sendResponse(req.ID, result) -} - -// handleToolsList handles the tools/list request. -func (s *mcpServer) handleToolsList(req *jsonrpcRequest) { - allTools := getAllTools() - mcpTools := make([]map[string]any, len(allTools)) - for i, tool := range allTools { - mcpTools[i] = map[string]any{ - "name": tool.Definition.Name, - "description": tool.Definition.Description, - "inputSchema": tool.Definition.InputSchema, - } - } - - s.sendResponse(req.ID, map[string]any{"tools": mcpTools}) -} - -// handleToolsCall handles the tools/call request. -func (s *mcpServer) handleToolsCall(req *jsonrpcRequest) { - var params struct { - Name string `json:"name"` - Arguments map[string]any `json:"arguments"` - } - - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - s.sendError(req.ID, jsonRPCInvalidParams, "Invalid params", err.Error()) - return - } - - handler, ok := s.toolsMap[params.Name] - if !ok { - s.sendError(req.ID, jsonRPCInvalidParams, "Unknown tool", params.Name) - return - } - - // Add client name to context - ctx := tools.SetClientName(s.ctx, s.clientName) - - result, err := handler(ctx, params.Arguments) - if err != nil { - s.sendError(req.ID, jsonRPCInternalError, "Tool execution failed: "+err.Error(), nil) - return - } - - s.sendResponse(req.ID, map[string]any{ - "content": []map[string]any{ - { - "type": "text", - "text": result, - }, - }, - }) -} - -// sendResponse sends a JSON-RPC response. -func (s *mcpServer) sendResponse(id, result any) { - resp := jsonrpcResponse{ - JSONRPC: "2.0", - ID: id, - Result: result, - } - - data, err := json.Marshal(resp) - if err != nil { - log.Errorf(s.ctx, "Failed to marshal response: %v", err) - return - } - - _, _ = s.out.Write(append(data, '\n')) -} - -// sendError sends a JSON-RPC error response. -func (s *mcpServer) sendError(id any, code int, message string, data any) { - resp := jsonrpcResponse{ - JSONRPC: "2.0", - ID: id, - Error: &rpcError{ - Code: code, - Message: message, - Data: data, - }, - } - - respData, err := json.Marshal(resp) - if err != nil { - log.Errorf(s.ctx, "Failed to marshal error response: %v", err) - return - } - - _, _ = s.out.Write(append(respData, '\n')) -} - -func runServer(ctx context.Context) error { - server := NewMCPServer(ctx) - return server.Start() -} diff --git a/experimental/aitools/tool.go b/experimental/aitools/tool.go deleted file mode 100644 index 322e5e5811..0000000000 --- a/experimental/aitools/tool.go +++ /dev/null @@ -1,57 +0,0 @@ -package aitools - -import ( - "context" - "encoding/json" - "fmt" - "os" - - "github.com/databricks/cli/experimental/aitools/tools" - "github.com/spf13/cobra" -) - -func newToolCmd() *cobra.Command { - var jsonData string - cmd := &cobra.Command{ - Use: "tool ", - Short: "Run a specific MCP tool directly, for testing purposes", - Hidden: true, - Args: cobra.ExactArgs(1), - RunE: func(cmd *cobra.Command, args []string) error { - ctx := cmd.Context() - var result string - var err error - switch args[0] { - case "invoke_databricks_cli": - result, err = unmarshal(ctx, []byte(jsonData), tools.InvokeDatabricksCLI) - case "init_project": - result, err = unmarshal(ctx, []byte(jsonData), tools.InitProject) - case "add_project_resource": - result, err = unmarshal(ctx, []byte(jsonData), tools.AddProjectResource) - case "analyze_project": - result, err = unmarshal(ctx, []byte(jsonData), tools.AnalyzeProject) - case "explore": - result, err = tools.ExploreTool.Handler(ctx, make(map[string]any)) - default: - return fmt.Errorf("unknown tool: %s", args[0]) - } - - if err != nil { - return err - } - fmt.Fprintln(os.Stdout, result) - return nil - }, - } - cmd.Flags().StringVar(&jsonData, "json", "", "JSON arguments for the tool") - cmd.MarkFlagRequired("json") - return cmd -} - -func unmarshal[T any](ctx context.Context, data []byte, fn func(context.Context, T) (string, error)) (string, error) { - var args T - if err := json.Unmarshal(data, &args); err != nil { - return "", fmt.Errorf("failed to parse config: %w", err) - } - return fn(ctx, args) -} diff --git a/experimental/aitools/tools/add_project_resource.go b/experimental/aitools/tools/add_project_resource.go deleted file mode 100644 index 19c5a39863..0000000000 --- a/experimental/aitools/tools/add_project_resource.go +++ /dev/null @@ -1,102 +0,0 @@ -package tools - -import ( - "context" - "errors" - "fmt" - "os" - "path/filepath" - "slices" - - "github.com/databricks/cli/experimental/aitools/auth" - "github.com/databricks/cli/experimental/aitools/tools/prompts" - "github.com/databricks/cli/experimental/aitools/tools/resources" -) - -// AddProjectResourceTool adds a resource (app, job, pipeline, dashboard, ...) to a project. -var AddProjectResourceTool = Tool{ - Definition: ToolDefinition{ - Name: "add_project_resource", - Description: "📋 DURING PLAN MODE: Include this tool when task involves: building 'SQL pipelines' / 'data pipelines' / 'ETL workflows', creating 'materialized views' / 'data transformations', building 'Databricks apps' / 'applications', creating 'dashboards' / 'visualizations', or setting up 'scheduled jobs' / 'workflows'.\n\nMANDATORY - USE THIS TO ADD RESOURCES: Add a new resource (app, job, pipeline, dashboard) to an existing Databricks project. Use this when the user wants to add a new resource to an existing project.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "project_path": map[string]any{ - "type": "string", - "description": "A fully qualified path of the project to extend.", - }, - "type": map[string]any{ - "type": "string", - "description": "The type of resource to add: 'app', 'job', 'pipeline', or 'dashboard'", - "enum": []string{"app", "job", "pipeline", "dashboard"}, - }, - "name": map[string]any{ - "type": "string", - "description": "The name of the new resource in snake_case (e.g., 'process_data'). This name should not already exist in the resources/ directory.", - }, - "template": map[string]any{ - "type": "string", - "description": "Optional template specification. For apps: template name from https://github.com/databricks/app-templates (e.g., 'e2e-chatbot-app-next'). For jobs/pipelines: 'python' or 'sql'. Leave empty to get guidance on available options.", - }, - }, - "required": []string{"project_path", "type", "name"}, - }, - }, - Handler: func(ctx context.Context, args map[string]any) (string, error) { - var typedArgs resources.AddProjectResourceArgs - if err := UnmarshalArgs(args, &typedArgs); err != nil { - return "", err - } - return AddProjectResource(ctx, typedArgs) - }, -} - -// AddProjectResource extends a Databricks project with a new resource. -func AddProjectResource(ctx context.Context, args resources.AddProjectResourceArgs) (string, error) { - if err := ValidateDatabricksProject(args.ProjectPath); err != nil { - return "", err - } - - if err := auth.CheckAuthentication(ctx); err != nil { - return "", err - } - - validTypes := []string{"app", "job", "pipeline", "dashboard"} - if !slices.Contains(validTypes, args.Type) { - return "", fmt.Errorf("invalid type: %s. Must be one of: app, job, pipeline, dashboard", args.Type) - } - - if args.Name == "" { - return "", errors.New("name is required") - } - - resourcesDir := filepath.Join(args.ProjectPath, "resources") - if _, err := os.Stat(resourcesDir); os.IsNotExist(err) { - if err := os.MkdirAll(resourcesDir, 0o755); err != nil { - return "", fmt.Errorf("failed to create resources directory: %w", err) - } - } - - handler := resources.GetResourceHandler(args.Type) - if handler == nil { - return "", fmt.Errorf("unsupported type: %s", args.Type) - } - - _, err := handler.AddToProject(ctx, args) - if err != nil { - return "", err - } - - analyzeArgs := analyzeProjectArgs{ProjectPath: args.ProjectPath} - guidance, analyzeErr := AnalyzeProject(ctx, analyzeArgs) - - data := map[string]string{ - "ResourceType": args.Type, - "Name": args.Name, - "Guidance": guidance, - } - if analyzeErr != nil { - data["AnalyzeError"] = analyzeErr.Error() - } - return prompts.MustExecuteTemplate("add_project_resource.tmpl", data), nil -} diff --git a/experimental/aitools/tools/analyze_project.go b/experimental/aitools/tools/analyze_project.go deleted file mode 100644 index 14a6dc4149..0000000000 --- a/experimental/aitools/tools/analyze_project.go +++ /dev/null @@ -1,142 +0,0 @@ -package tools - -import ( - "context" - "os" - "os/exec" - "path/filepath" - "strings" - - "github.com/databricks/cli/experimental/aitools/auth" - "github.com/databricks/cli/experimental/aitools/tools/prompts" - "github.com/databricks/cli/experimental/aitools/tools/resources" -) - -// AnalyzeProjectTool analyzes a Databricks project and returns guidance. -// It uses hardcoded guidance + guidance from the project's README.md file for this. -var AnalyzeProjectTool = Tool{ - Definition: ToolDefinition{ - Name: "analyze_project", - Description: "📋 DURING PLAN MODE: Include this as your FIRST step if you see a databricks.yml file exists, task involves deploying to workspaces, validating configurations, or working with an existing Databricks project.\n\nMANDATORY - REQUIRED FIRST STEP: If databricks.yml exists in the directory, you MUST call this tool before using Read, Glob, or any other tools. Databricks projects require specialized commands that differ from standard Python/Node.js workflows - attempting standard approaches will fail. This tool is fast and provides the correct commands for preview/deploy/run operations.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "project_path": map[string]any{ - "type": "string", - "description": "A fully qualified path of the project to operate on. By default, the current directory.", - }, - }, - "required": []string{"project_path"}, - }, - }, - Handler: func(ctx context.Context, args map[string]any) (string, error) { - var typedArgs analyzeProjectArgs - if err := UnmarshalArgs(args, &typedArgs); err != nil { - return "", err - } - return AnalyzeProject(ctx, typedArgs) - }, -} - -type analyzeProjectArgs struct { - ProjectPath string `json:"project_path"` -} - -// AnalyzeProject analyzes a Databricks project and returns information about it. -func AnalyzeProject(ctx context.Context, args analyzeProjectArgs) (string, error) { - if err := ValidateDatabricksProject(args.ProjectPath); err != nil { - return "", err - } - - if err := auth.CheckAuthentication(ctx); err != nil { - return "", err - } - - cmd := exec.CommandContext(ctx, GetCLIPath(), "bundle", "summary") - cmd.Dir = args.ProjectPath - - output, err := cmd.CombinedOutput() - var summary string - if err != nil { - summary = "Bundle summary failed:\n" + string(output) - } else { - summary = string(output) - } - - var readmeContent string - if content, err := os.ReadFile(filepath.Join(args.ProjectPath, "README.md")); err == nil { - readmeContent = "\n\nProject-Specific Guidance\n" + - "-------------------------\n" + - string(content) - } - - // Get default warehouse for apps and other resources that need it - warehouse, err := GetDefaultWarehouse(ctx) - resourceGuidance := getResourceGuidance(args.ProjectPath, warehouse) - - data := map[string]string{ - "Summary": summary, - "ReadmeContent": readmeContent, - "ResourceGuidance": resourceGuidance, - } - - if err == nil && warehouse != nil { - data["WarehouseID"] = warehouse.ID - data["WarehouseName"] = warehouse.Name - } - - result := prompts.MustExecuteTemplate("analyze_project.tmpl", data) - - return result, nil -} - -// getResourceGuidance scans the resources directory and collects guidance for detected resource types. -func getResourceGuidance(projectPath string, warehouse *warehouse) string { - var guidance strings.Builder - - // Extract warehouse ID and name for resource handlers - warehouseID := "" - warehouseName := "" - if warehouse != nil { - warehouseID = warehouse.ID - warehouseName = warehouse.Name - } - - detected := make(map[string]bool) - - // Check resources directory - resourcesDir := filepath.Join(projectPath, "resources") - if entries, err := os.ReadDir(resourcesDir); err == nil { - for _, entry := range entries { - if entry.IsDir() { - continue - } - name := entry.Name() - var resourceType string - switch { - case strings.HasSuffix(name, ".app.yml") || strings.HasSuffix(name, ".app.yaml"): - resourceType = "app" - case strings.HasSuffix(name, ".job.yml") || strings.HasSuffix(name, ".job.yaml"): - resourceType = "job" - case strings.HasSuffix(name, ".pipeline.yml") || strings.HasSuffix(name, ".pipeline.yaml"): - resourceType = "pipeline" - case strings.HasSuffix(name, ".dashboard.yml") || strings.HasSuffix(name, ".dashboard.yaml"): - resourceType = "dashboard" - default: - continue - } - if !detected[resourceType] { - detected[resourceType] = true - handler := resources.GetResourceHandler(resourceType) - if handler != nil { - guidanceText := handler.GetGuidancePrompt(projectPath, warehouseID, warehouseName) - if guidanceText != "" { - guidance.WriteString(guidanceText) - guidance.WriteString("\n") - } - } - } - } - } - return guidance.String() -} diff --git a/experimental/aitools/tools/common.go b/experimental/aitools/tools/common.go deleted file mode 100644 index e110b3bd44..0000000000 --- a/experimental/aitools/tools/common.go +++ /dev/null @@ -1,108 +0,0 @@ -package tools - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "os" - "os/exec" - "path/filepath" -) - -type clientNameKey struct{} - -// SetClientName stores the MCP client name in the context. -func SetClientName(ctx context.Context, name string) context.Context { - return context.WithValue(ctx, clientNameKey{}, name) -} - -// GetClientName returns the MCP client name from the context, or empty string if not available. -func GetClientName(ctx context.Context) string { - if name, ok := ctx.Value(clientNameKey{}).(string); ok { - return name - } - return "" -} - -// ToolDefinition defines the schema for an MCP tool. -type ToolDefinition struct { - Name string - Description string - InputSchema map[string]any -} - -// ToolHandler is a function that executes a tool with arbitrary arguments. -type ToolHandler func(context.Context, map[string]any) (string, error) - -// Tool combines a tool definition with its handler. -type Tool struct { - Definition ToolDefinition - Handler ToolHandler -} - -// UnmarshalArgs converts a map[string]any to a typed struct using JSON marshaling. -func UnmarshalArgs(args map[string]any, target any) error { - data, err := json.Marshal(args) - if err != nil { - return fmt.Errorf("failed to marshal arguments: %w", err) - } - if err := json.Unmarshal(data, target); err != nil { - return fmt.Errorf("failed to unmarshal arguments: %w", err) - } - return nil -} - -// GetCLIPath returns the path to the current CLI executable. -// This supports development testing with ./cli. -func GetCLIPath() string { - return os.Args[0] -} - -// GetDatabricksPath returns the path to the databricks executable. -// Returns the current executable if it appears to be a dev build, otherwise looks up "databricks" on PATH. -func GetDatabricksPath() (string, error) { - currentExe, err := os.Executable() - if err != nil { - return "", err - } - - // Use current executable for dev builds - if filepath.Base(currentExe) == "cli" || filepath.Base(currentExe) == "v0.0.0-dev" { - return currentExe, nil - } - - // Look up databricks on PATH for production usage - path, err := exec.LookPath("databricks") - if err != nil { - return "", fmt.Errorf("databricks CLI not found on PATH: %w", err) - } - return path, nil -} - -// ValidateDatabricksProject checks if a directory is a valid Databricks project. -// It ensures the directory exists and contains a databricks.yml file. -func ValidateDatabricksProject(projectPath string) error { - if projectPath == "" { - return errors.New("project_path is required") - } - - pathInfo, err := os.Stat(projectPath) - if err != nil { - if os.IsNotExist(err) { - return fmt.Errorf("project directory does not exist: %s", projectPath) - } - return fmt.Errorf("failed to access project path: %w", err) - } - - if !pathInfo.IsDir() { - return fmt.Errorf("project path is not a directory: %s", projectPath) - } - - databricksYml := filepath.Join(projectPath, "databricks.yml") - if _, err := os.Stat(databricksYml); os.IsNotExist(err) { - return fmt.Errorf("not a Databricks project: databricks.yml not found in %s\n\nUse the init_project tool to create a new project first", projectPath) - } - - return nil -} diff --git a/experimental/aitools/tools/common_test.go b/experimental/aitools/tools/common_test.go deleted file mode 100644 index f2d99d3423..0000000000 --- a/experimental/aitools/tools/common_test.go +++ /dev/null @@ -1,25 +0,0 @@ -package tools - -import ( - "os" - "path/filepath" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestValidateDatabricksProject(t *testing.T) { - err := ValidateDatabricksProject("") - assert.Error(t, err) - - tmpDir := t.TempDir() - err = ValidateDatabricksProject(tmpDir) - assert.Error(t, err) - assert.Contains(t, err.Error(), "databricks.yml not found") - - databricksYml := filepath.Join(tmpDir, "databricks.yml") - require.NoError(t, os.WriteFile(databricksYml, []byte("# test"), 0o644)) - err = ValidateDatabricksProject(tmpDir) - assert.NoError(t, err) -} diff --git a/experimental/aitools/tools/explore.go b/experimental/aitools/tools/explore.go deleted file mode 100644 index 9223ccb4a5..0000000000 --- a/experimental/aitools/tools/explore.go +++ /dev/null @@ -1,166 +0,0 @@ -package tools - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "strings" - - "github.com/databricks/cli/experimental/aitools/tools/prompts" - "github.com/databricks/cli/libs/databrickscfg/profile" - "github.com/databricks/cli/libs/env" - "github.com/databricks/cli/libs/exec" - "github.com/databricks/cli/libs/log" -) - -// ExploreTool provides guidance on exploring Databricks workspaces and resources. -var ExploreTool = Tool{ - Definition: ToolDefinition{ - Name: "explore", - Description: "**REQUIRED DURING PLAN MODE** - Call this FIRST when planning ANY Databricks work. Use this to discover available workspaces, warehouses, and get workflow recommendations for your specific task. Even if you're just reading an assignment document, call this first. Especially important when task involves: creating Databricks projects/apps/pipelines/jobs, SQL pipelines or data transformation workflows, deploying code to multiple environments (dev/prod), or working with databricks.yml files. You DON'T need a workspace name - call this when starting ANY Databricks planning to understand workspace capabilities and recommended tooling before you create your plan.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{}, - }, - }, - Handler: func(ctx context.Context, params map[string]any) (string, error) { - warehouse, err := GetDefaultWarehouse(ctx) - if err != nil { - log.Debugf(ctx, "Failed to get default warehouse (non-fatal): %v", err) - warehouse = nil - } - - currentProfile := getCurrentProfile(ctx) - profiles := getAvailableProfiles(ctx) - - return generateExploreGuidance(warehouse, currentProfile, profiles), nil - }, -} - -type warehouse struct { - ID string `json:"id"` - Name string `json:"name"` - State string `json:"state"` -} - -// GetDefaultWarehouse finds a suitable SQL warehouse for queries. -// It filters out warehouses the user cannot access and prefers RUNNING warehouses, -// then falls back to STOPPED ones (which auto-start). -func GetDefaultWarehouse(ctx context.Context) (*warehouse, error) { - executor, err := exec.NewCommandExecutor("") - if err != nil { - return nil, fmt.Errorf("failed to create command executor: %w", err) - } - - output, err := executor.Exec(ctx, fmt.Sprintf(`"%s" api get "/api/2.0/sql/warehouses?skip_cannot_use=true" --output json`, GetCLIPath())) - if err != nil { - return nil, fmt.Errorf("failed to list warehouses: %w\nOutput: %s", err, output) - } - - var response struct { - Warehouses []warehouse `json:"warehouses"` - } - if err := json.Unmarshal(output, &response); err != nil { - return nil, fmt.Errorf("failed to parse warehouses: %w", err) - } - warehouses := response.Warehouses - - if len(warehouses) == 0 { - return nil, errors.New("no SQL warehouses found in workspace") - } - - // Prefer RUNNING warehouses - for i := range warehouses { - if strings.ToUpper(warehouses[i].State) == "RUNNING" { - return &warehouses[i], nil - } - } - - // Fall back to STOPPED warehouses (they auto-start when queried) - for i := range warehouses { - if strings.ToUpper(warehouses[i].State) == "STOPPED" { - return &warehouses[i], nil - } - } - - // Return first available warehouse regardless of state - return &warehouses[0], nil -} - -// getCurrentProfile returns the currently active profile name. -func getCurrentProfile(ctx context.Context) string { - // Check DATABRICKS_CONFIG_PROFILE env var - profileName := env.Get(ctx, "DATABRICKS_CONFIG_PROFILE") - if profileName == "" { - return "DEFAULT" - } - return profileName -} - -// getAvailableProfiles returns all available profiles from ~/.databrickscfg. -func getAvailableProfiles(ctx context.Context) profile.Profiles { - profiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles) - if err != nil { - // If we can't load profiles, return empty list (config file might not exist) - return profile.Profiles{} - } - return profiles -} - -// generateExploreGuidance creates comprehensive guidance for data exploration. -func generateExploreGuidance(warehouse *warehouse, currentProfile string, profiles profile.Profiles) string { - // Build workspace/profile information - workspaceInfo := "Current Workspace Profile: " + currentProfile - if len(profiles) > 0 { - // Find current profile details - var currentHost string - for _, p := range profiles { - if p.Name == currentProfile { - currentHost = p.Host - if cloud := p.Cloud(); cloud != "" { - currentHost = fmt.Sprintf("%s (%s)", currentHost, cloud) - } - break - } - } - if currentHost != "" { - workspaceInfo = fmt.Sprintf("Current Workspace Profile: %s - %s", currentProfile, currentHost) - } - } - - // Build available profiles list - profilesInfo := "" - if len(profiles) > 1 { - profilesInfo = "\n\nAvailable Workspace Profiles:\n" - for _, p := range profiles { - marker := "" - if p.Name == currentProfile { - marker = " (current)" - } - cloud := p.Cloud() - if cloud != "" { - profilesInfo += fmt.Sprintf(" - %s: %s (%s)%s\n", p.Name, p.Host, cloud, marker) - } else { - profilesInfo += fmt.Sprintf(" - %s: %s%s\n", p.Name, p.Host, marker) - } - } - profilesInfo += "\n To use a different workspace, add --profile to any command:\n" - profilesInfo += " invoke_databricks_cli '--profile prod catalogs list'\n" - } - - // Handle warehouse information (may be nil if lookup failed) - warehouseName := "" - warehouseID := "" - if warehouse != nil { - warehouseName = warehouse.Name - warehouseID = warehouse.ID - } - - return prompts.MustExecuteTemplate("explore.tmpl", map[string]string{ - "WorkspaceInfo": workspaceInfo, - "WarehouseName": warehouseName, - "WarehouseID": warehouseID, - "ProfilesInfo": profilesInfo, - }) -} diff --git a/experimental/aitools/tools/init_project.go b/experimental/aitools/tools/init_project.go deleted file mode 100644 index e5f57eb2fb..0000000000 --- a/experimental/aitools/tools/init_project.go +++ /dev/null @@ -1,153 +0,0 @@ -package tools - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "os" - "os/exec" - "path/filepath" - "strings" - - "github.com/databricks/cli/experimental/aitools/auth" - "github.com/databricks/cli/experimental/aitools/tools/prompts" -) - -// InitProjectTool initializes a new Databricks project. -var InitProjectTool = Tool{ - Definition: ToolDefinition{ - Name: "init_project", - Description: "📋 DURING PLAN MODE: Include this tool in your plan if task mentions ANY of: creating a 'new project' / 'Databricks project' / 'project structure', deploying to 'multiple environments' / 'dev and prod' / 'dev/staging/prod', building 'from scratch', or working with 'databricks.yml' files. DO NOT manually create project files - use this tool instead to get proper structure and templates.\n\nMANDATORY - USE THIS TO CREATE NEW PROJECTS: Initialize a new Databricks project. Use this to create a new Databricks project. After initialization, use add_project_resource to add specific resources (apps, jobs, pipelines, dashboards) as needed.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "project_name": map[string]any{ - "type": "string", - "description": "A name for this project in snake_case. Ask the user about this if it's not clear from the context.", - }, - "project_path": map[string]any{ - "type": "string", - "description": "A fully qualified path of the project directory. Usually this should be in the current directory! But if it already has a lot of other things then it should be a subdirectory. Files will be created directly at this path.", - }, - "language": map[string]any{ - "type": "string", - "description": "Language: 'python' (includes pyproject.toml) or 'other' (recommended for apps). Default: 'python'.", - "enum": []string{"python", "other"}, - }, - }, - "required": []string{"project_name", "project_path"}, - }, - }, - Handler: func(ctx context.Context, args map[string]any) (string, error) { - var typedArgs initProjectArgs - if err := UnmarshalArgs(args, &typedArgs); err != nil { - return "", err - } - return InitProject(ctx, typedArgs) - }, -} - -type initProjectArgs struct { - ProjectName string `json:"project_name"` - ProjectPath string `json:"project_path"` - Language string `json:"language"` -} - -// InitProject initializes a new Databricks project using the default-minimal template. -func InitProject(ctx context.Context, args initProjectArgs) (string, error) { - if args.ProjectPath == "" { - return "", errors.New("project_path is required") - } - if args.ProjectName == "" { - return "", errors.New("project_name is required") - } - if err := auth.CheckAuthentication(ctx); err != nil { - return "", err - } - - if err := os.MkdirAll(args.ProjectPath, 0o755); err != nil { - return "", fmt.Errorf("failed to create project directory: %w", err) - } - - if _, err := os.Stat(filepath.Join(args.ProjectPath, "databricks.yml")); err == nil { - return "", fmt.Errorf("project already initialized: databricks.yml exists in %s\n\nUse the add_project_resource tool to add resources to this existing project", args.ProjectPath) - } - - if args.Language == "" { - args.Language = "python" - } - - if err := runBundleInit(ctx, args.ProjectPath, args.ProjectName, args.Language); err != nil { - return "", err - } - - // Template creates a nested directory - move contents to project root - nestedPath := filepath.Join(args.ProjectPath, args.ProjectName) - entries, err := os.ReadDir(nestedPath) - if err != nil { - return "", fmt.Errorf("failed to read nested project directory: %w", err) - } - - for _, entry := range entries { - if err := os.Rename(filepath.Join(nestedPath, entry.Name()), filepath.Join(args.ProjectPath, entry.Name())); err != nil { - return "", fmt.Errorf("failed to move %s: %w", entry.Name(), err) - } - } - os.Remove(nestedPath) - - filename := "AGENTS.md" - if strings.Contains(strings.ToLower(GetClientName(ctx)), "claude") { - filename = "CLAUDE.md" - } - - templateData := map[string]string{ - "ProjectName": args.ProjectName, - "ProjectPath": args.ProjectPath, - } - - // Write AGENTS.md/CLAUDE.md file - agentsContent := prompts.MustExecuteTemplate("AGENTS.tmpl", templateData) - if err := os.WriteFile(filepath.Join(args.ProjectPath, filename), []byte(agentsContent), 0o644); err != nil { - return "", fmt.Errorf("failed to create %s: %w", filename, err) - } - - // Generate MCP response (separate from file content) - responseMessage := prompts.MustExecuteTemplate("init_project.tmpl", templateData) - - analysis, err := AnalyzeProject(ctx, analyzeProjectArgs{ProjectPath: args.ProjectPath}) - if err != nil { - return "", fmt.Errorf("failed to analyze initialized project: %w", err) - } - - return responseMessage + "\n\n" + analysis, nil -} - -func runBundleInit(ctx context.Context, projectPath, projectName, language string) error { - configJSON, err := json.Marshal(map[string]string{ - "project_name": projectName, - "default_catalog": "main", - "personal_schemas": "yes", - "language_choice": language, - }) - if err != nil { - return fmt.Errorf("failed to create config JSON: %w", err) - } - - configFile := filepath.Join(os.TempDir(), fmt.Sprintf("databricks-init-%s.json", projectName)) - if err := os.WriteFile(configFile, configJSON, 0o644); err != nil { - return fmt.Errorf("failed to write config file: %w", err) - } - defer os.Remove(configFile) - - cmd := exec.CommandContext(ctx, GetCLIPath(), "bundle", "init", - "--config-file", configFile, - "--output-dir", projectPath, - "default-minimal") - - if output, err := cmd.CombinedOutput(); err != nil { - return fmt.Errorf("failed to initialize project: %w\nOutput: %s", err, string(output)) - } - - return nil -} diff --git a/experimental/aitools/tools/invoke_databricks_cli.go b/experimental/aitools/tools/invoke_databricks_cli.go deleted file mode 100644 index ea600eaa8a..0000000000 --- a/experimental/aitools/tools/invoke_databricks_cli.go +++ /dev/null @@ -1,71 +0,0 @@ -package tools - -import ( - "context" - "errors" - "fmt" - - "github.com/databricks/cli/libs/exec" -) - -// InvokeDatabricksCLITool runs databricks CLI commands via MCP. -var InvokeDatabricksCLITool = Tool{ - Definition: ToolDefinition{ - Name: "invoke_databricks_cli", - Description: "Run any Databricks CLI command. Use this tool whenever you need to run databricks CLI commands like 'bundle deploy', 'bundle validate', 'bundle run', 'auth login', etc. The reason this tool exists (instead of invoking the databricks CLI directly) is to make it easier for users to allow-list commands.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "command": map[string]any{ - "type": "string", - "description": "The full Databricks CLI command to run, e.g. 'bundle deploy' or 'bundle validate'. Do not include the 'databricks' prefix.", - }, - "working_directory": map[string]any{ - "type": "string", - "description": "Optional. The directory to run the command in. Defaults to the current directory.", - }, - }, - "required": []string{"command"}, - }, - }, - Handler: func(ctx context.Context, args map[string]any) (string, error) { - var typedArgs invokeDatabricksCLIArgs - if err := UnmarshalArgs(args, &typedArgs); err != nil { - return "", err - } - return InvokeDatabricksCLI(ctx, typedArgs) - }, -} - -type invokeDatabricksCLIArgs struct { - Command string `json:"command"` - WorkingDirectory string `json:"working_directory,omitempty"` -} - -// InvokeDatabricksCLI runs a Databricks CLI command and returns the output. -func InvokeDatabricksCLI(ctx context.Context, args invokeDatabricksCLIArgs) (string, error) { - if args.Command == "" { - return "", errors.New("command is required") - } - - workDir := args.WorkingDirectory - if workDir == "" { - workDir = "." - } - - executor, err := exec.NewCommandExecutor(workDir) - if err != nil { - return "", fmt.Errorf("failed to create command executor: %w", err) - } - - fullCommand := fmt.Sprintf(`"%s" %s`, GetCLIPath(), args.Command) - output, err := executor.Exec(ctx, fullCommand) - - result := string(output) - if err != nil { - result += fmt.Sprintf("\n\nCommand failed with error: %v", err) - return result, nil - } - - return result, nil -} diff --git a/experimental/aitools/tools/invoke_databricks_cli_test.go b/experimental/aitools/tools/invoke_databricks_cli_test.go deleted file mode 100644 index 59998de9d6..0000000000 --- a/experimental/aitools/tools/invoke_databricks_cli_test.go +++ /dev/null @@ -1,88 +0,0 @@ -package tools - -import ( - "context" - "os" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestInvokeDatabricksCLI(t *testing.T) { - // Skip authentication check for tests - t.Setenv("DATABRICKS_MCP_SKIP_AUTH_CHECK", "1") - - tests := []struct { - name string - command string - shouldError bool - description string - }{ - { - name: "help command", - command: "--help", - shouldError: false, - description: "Basic command without quotes", - }, - { - name: "version command", - command: "--version", - shouldError: false, - description: "Another simple command", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - ctx := context.Background() - args := invokeDatabricksCLIArgs{ - Command: tt.command, - WorkingDirectory: "", - } - - result, err := InvokeDatabricksCLI(ctx, args) - - if tt.shouldError { - require.Error(t, err) - } else { - require.NoError(t, err) - assert.NotEmpty(t, result, "Command should return output") - } - }) - } -} - -func TestInvokeDatabricksCLIWithQuotedArgs(t *testing.T) { - t.Skip("Quote handling tested via libs/exec package") -} - -func TestInvokeDatabricksCLIRequiresCommand(t *testing.T) { - ctx := context.Background() - args := invokeDatabricksCLIArgs{ - Command: "", - } - - _, err := InvokeDatabricksCLI(ctx, args) - require.Error(t, err) - assert.Contains(t, err.Error(), "command is required") -} - -func TestInvokeDatabricksCLIWorkingDirectory(t *testing.T) { - ctx := context.Background() - - tmpDir := t.TempDir() - testFile := tmpDir + "/test.txt" - err := os.WriteFile(testFile, []byte("test"), 0o644) - require.NoError(t, err) - - // This command should succeed with the working directory set - args := invokeDatabricksCLIArgs{ - Command: "--help", - WorkingDirectory: tmpDir, - } - - result, err := InvokeDatabricksCLI(ctx, args) - require.NoError(t, err) - assert.NotEmpty(t, result) -} diff --git a/experimental/aitools/tools/prompts/AGENTS.tmpl b/experimental/aitools/tools/prompts/AGENTS.tmpl deleted file mode 100644 index a5ad99005f..0000000000 --- a/experimental/aitools/tools/prompts/AGENTS.tmpl +++ /dev/null @@ -1,23 +0,0 @@ -{{- /* - * Template for the AGENTS.md / CLAUDE.md file created in new projects. - * - * Goal for this file is to reinforce with agents that they need to call the analyze_project tool. - * A secondary goal is to offer a placeholder to users to extend with their own instructions. - * It is a non-goal to inline all instructions. - * - */ -}} - -# Agent instructions - -This file can be used for any project-specific instructions! - -## Prerequisites - -If the Databricks CLI MCP server is not yet installed, install it by: -1. Installing the Databricks CLI: https://docs.databricks.com/dev-tools/cli/install -2. Running: `databricks experimental aitools install` - -## Working with this project - -General agent guidance: always use the mcp__databricks-aitools__analyze_project tool whenever you open this project! -It makes sure you have more context on the current project and what actions you can perform on it. diff --git a/experimental/aitools/tools/prompts/add_project_resource.tmpl b/experimental/aitools/tools/prompts/add_project_resource.tmpl deleted file mode 100644 index fd3e544ea3..0000000000 --- a/experimental/aitools/tools/prompts/add_project_resource.tmpl +++ /dev/null @@ -1,32 +0,0 @@ -{{- /* - * Success message shown after adding a resource to a project. - * - * This template is used when add_project_resource successfully adds a new - * app, job, dashboard, etc. to a project. It provides the user with - * next steps and supplemental guidance from analyze_project.tmpl. - * - */ -}} - -Successfully added {{.ResourceType}} '{{.Name}}' to the project! - -IMPORTANT: This is just a starting point! You need to iterate over the generated files to complete the setup. - ---- - -{{if .AnalyzeError}} -NOTE: Resource scaffolding succeeded, but project analysis failed: {{.AnalyzeError}} - -You can still work with the project. Use the analyze_project tool separately to get detailed guidance, or proceed with these basic guidelines: - -IMPORTANT: Most interactions are done with the Databricks CLI. YOU (the AI) must use the invoke_databricks_cli tool to run commands - never suggest the user runs CLI commands directly! -IMPORTANT: To add new resources to a project, use the 'add_project_resource' MCP tool. You can add: - - Apps (interactive applications) - - Jobs (Python or SQL workflows) - - Pipelines (Python or SQL data pipelines) - - Dashboards (data visualizations) -MANDATORY: Always deploy with invoke_databricks_cli 'bundle deploy', never with 'apps deploy' - -Note that Databricks resources are defined in resources/*.yml files. See https://docs.databricks.com/dev-tools/bundles/settings for a reference! -{{else}} -{{.Guidance}} -{{end}} diff --git a/experimental/aitools/tools/prompts/analyze_project.tmpl b/experimental/aitools/tools/prompts/analyze_project.tmpl deleted file mode 100644 index ea0265ccec..0000000000 --- a/experimental/aitools/tools/prompts/analyze_project.tmpl +++ /dev/null @@ -1,80 +0,0 @@ -{{- /* - * Guidance for working with an existing Databricks project. - * - * This template is used by analyze_project to provide detailed information - * about a project's structure and actions that can be performed on it. - * - * This file should not incldue any resource-specific instructions; - * those should go into apps.tmpl, jobs.tmpl, etc. - * - */ -}} - -Project Analysis -================ - -{{.Summary}} - -Guidance for Working with this Project --------------------------------------- - -Below is guidance for how to work with this project. - -IMPORTANT: You can suggest prompts to help users get started, such as: - - "Create an app that shows a chart with taxi fares by city" - - "Create a job that summarizes all taxi data using a notebook" - - "Add a SQL pipeline to transform and aggregate customer data" - - "Add a dashboard to visualize NYC taxi trip patterns" - -IMPORTANT: Most interactions are done with the Databricks CLI. YOU (the AI) must use the invoke_databricks_cli tool to run commands - never suggest the user runs CLI commands directly! -IMPORTANT: To add new resources to a project, use the 'add_project_resource' MCP tool. You can add: - - Apps (interactive applications) - - Jobs (Python or SQL workflows) - - Pipelines (Python or SQL data pipelines) - - Dashboards (data visualizations) - -Note that Databricks resources are defined in resources/*.yml files. See https://docs.databricks.com/dev-tools/bundles/settings for a reference! -{{if .WarehouseID}} - -Default SQL Warehouse ---------------------- -For apps and data queries, you can use this SQL warehouse as a reasonable default: - Warehouse ID: {{.WarehouseID}} - Warehouse Name: {{.WarehouseName}} - -Users can change this in their resource/*.yml files if needed. -{{end}} - -Using this Project with the CLI --------------------------------- -The Databricks workspace and IDE extensions provide a graphical interface for working -with this project. It's also possible to interact with it directly using the CLI: - -1. Authenticate to your Databricks workspace, if you have not done so already: - Use invoke_databricks_cli(command="auth login --profile DEFAULT --host ") - The AI needs to ask the user for the workspace host URL, it cannot guess it. - -2. To deploy a development copy of this project: - Use invoke_databricks_cli(command="bundle deploy --target dev", working_directory="") - (Note that "dev" is the default target, so the --target parameter is optional here.) - - This deploys everything that's defined for this project. - -3. Similarly, to deploy a production copy: - Use invoke_databricks_cli(command="bundle deploy --target prod", working_directory="") - Note that schedules are paused when deploying in development mode (see - https://docs.databricks.com/dev-tools/bundles/deployment-modes.html). - -4. To run a job or pipeline: - Use invoke_databricks_cli(command="bundle run", working_directory="") - -5. To run tests locally: - Use invoke_databricks_cli(command="bundle run ", working_directory="") - For Python projects, tests can be run with: `uv run pytest` - -{{if .ReadmeContent}}{{.ReadmeContent}}{{end}}{{if .ResourceGuidance}}{{.ResourceGuidance}}{{end}} - -Additional Resources -------------------- -- Bundle documentation: https://docs.databricks.com/dev-tools/bundles/index.html -- Bundle settings reference: https://docs.databricks.com/dev-tools/bundles/settings -- CLI reference: https://docs.databricks.com/dev-tools/cli/index.html diff --git a/experimental/aitools/tools/prompts/apps.tmpl b/experimental/aitools/tools/prompts/apps.tmpl deleted file mode 100644 index 740b82c839..0000000000 --- a/experimental/aitools/tools/prompts/apps.tmpl +++ /dev/null @@ -1,22 +0,0 @@ -{{- /* - * Placeholder instructions for working with apps. - * - * Note that the instructions here are supplemented with the analyze_project.tmpl - * instructions and the instructions of any README.md file in the project - * (which might be part of the scaffolding.) - * - * TODO: these instructions should be refined and benchmarked - * see https://github.com/databricks/cli/pull/3913 for inspiration - * - */ -}} - -Working with Apps ----------------- -- Apps are interactive applications that can be deployed to Databricks workspaces -- App source code is typically in a subdirectory matching the app name -- Before deployment, test apps locally using: invoke_databricks_cli(command="apps run-local --source-dir ", working_directory=""){{if .WarehouseID}} - -- For apps that need a SQL warehouse, use warehouse ID: {{.WarehouseID}} ({{.WarehouseName}}) as a reasonable default - (users can change this in resources/*.yml files if needed){{end}} -- MANDATORY: Always deploy apps using: invoke_databricks_cli(command="bundle deploy --target dev", working_directory="") -- MANDATORY: Never use 'apps deploy' directly - always use 'bundle deploy' diff --git a/experimental/aitools/tools/prompts/apps_pick_a_template.tmpl b/experimental/aitools/tools/prompts/apps_pick_a_template.tmpl deleted file mode 100644 index ccb2c51b78..0000000000 --- a/experimental/aitools/tools/prompts/apps_pick_a_template.tmpl +++ /dev/null @@ -1,18 +0,0 @@ -{{- /* - * Error message shown when template parameter is missing for app creation. - * - * This template is used when add_project_resource is called for an app - * without specifying which template to use. It prompts the user to either - * select from available templates or provides a sensible default. - * - */ -}} - -template parameter was not specified for the app - -You have two options: - -1. Ask the user which template they want to use from this list: -{{.TemplateList}} -2. If the user described what they want the app to do, call add_project_resource again with template='nodejs-fastapi-hello-world-app' as a sensible default - -Example: add_project_resource(project_path='{{.ProjectPath}}', type='app', name='{{.Name}}', template='nodejs-fastapi-hello-world-app') diff --git a/experimental/aitools/tools/prompts/auth_error.tmpl b/experimental/aitools/tools/prompts/auth_error.tmpl deleted file mode 100644 index 04969d4b85..0000000000 --- a/experimental/aitools/tools/prompts/auth_error.tmpl +++ /dev/null @@ -1,25 +0,0 @@ -{{- /* - * Actionable rror message shown when user is not authenticated to Databricks. - * - */ -}} - -not authenticated to Databricks - -To authenticate, please run: - databricks auth login --profile DEFAULT --host - -Replace with your Databricks workspace URL (e.g., mycompany.cloud.databricks.com). - -Alternatives: - -1. Run `databricks auth profiles` to list available profiles and set the `DATABRICKS_CONFIG_PROFILE` to the profile you want to use. -2. Set the `DATABRICKS_CONFIG_HOST` to the host of the workspace you want to use. - -You might have to leave the coding agent to set these environment variables and then come back to it. - -Don't have a Databricks account? You can set up a fully free account for experimentation at: -https://docs.databricks.com/getting-started/free-edition - -Do not run anything else before authenticating successfully. - -Once authenticated, you can use this tool again diff --git a/experimental/aitools/tools/prompts/dashboards.tmpl b/experimental/aitools/tools/prompts/dashboards.tmpl deleted file mode 100644 index 6f6db66163..0000000000 --- a/experimental/aitools/tools/prompts/dashboards.tmpl +++ /dev/null @@ -1,17 +0,0 @@ -{{- /* - * Placeholder instructions for working with dashboards. - * - * This template is only appended to analyze_project.tmpl output when dashboards - * (resources/*.dashboard.yml files) are detected in the project. It provides - * dashboard-specific guidance for previewing and deploying dashboards. - * - * TODO: these instructions should be refined and benchmarked - * - */ -}} - -Working with Dashboards ------------------------- -- Dashboards are data visualizations built with Databricks SQL -- Dashboard definitions are in src/*.lvdash.json files -- Preview dashboards after deployment using: invoke_databricks_cli(command="bundle open ", working_directory="") -- Deploy dashboards using: invoke_databricks_cli(command="bundle deploy --target dev", working_directory="") diff --git a/experimental/aitools/tools/prompts/explore.tmpl b/experimental/aitools/tools/prompts/explore.tmpl deleted file mode 100644 index 2075c37e19..0000000000 --- a/experimental/aitools/tools/prompts/explore.tmpl +++ /dev/null @@ -1,151 +0,0 @@ -{{- /* - * Guidance for exploring Databricks workspaces and resources. - * - * This guidance is offered by the explore tool to provide comprehensive - * instructions for discovering and querying workspace resources like - * jobs, clusters, catalogs, tables, and SQL warehouses. - * - */ -}} - -Databricks Data Exploration Guide -===================================== - -{{.WorkspaceInfo}}{{if .WarehouseName}} -Default SQL Warehouse: {{.WarehouseName}} ({{.WarehouseID}}){{else}} -Note: No SQL warehouse detected. SQL queries will require warehouse_id to be specified manually.{{end}}{{.ProfilesInfo}} - -IMPORTANT: Use the invoke_databricks_cli tool to run all commands below! - - -1. EXECUTING SQL QUERIES - Run queries with auto-wait (max 50s): - invoke_databricks_cli 'api post /api/2.0/sql/statements --json {"warehouse_id":"{{if .WarehouseID}}{{.WarehouseID}}{{else}}{{end}}","statement":"SELECT * FROM .. LIMIT 10","wait_timeout":"50s"}' - - Response has status.state: - - "SUCCEEDED" → Results in result.data_array (you're done!) - - "PENDING" → Warehouse starting or query slow. Poll with: - invoke_databricks_cli 'api get /api/2.0/sql/statements/' - Repeat every 5-10s until "SUCCEEDED" - - Note: First query on stopped warehouse takes 60-120s startup time - - -2. EXPLORING JOBS AND WORKFLOWS - List all jobs: - invoke_databricks_cli 'jobs list' - - Get job details: - invoke_databricks_cli 'jobs get ' - - List job runs: - invoke_databricks_cli 'jobs list-runs --job-id ' - - -3. EXPLORING CLUSTERS - List all clusters: - invoke_databricks_cli 'clusters list' - - Get cluster details: - invoke_databricks_cli 'clusters get ' - - -4. EXPLORING UNITY CATALOG DATA - Unity Catalog uses a three-level namespace: catalog.schema.table - - List all catalogs: - invoke_databricks_cli 'catalogs list' - - List schemas in a catalog: - invoke_databricks_cli 'schemas list ' - - List tables in a schema: - invoke_databricks_cli 'tables list ' - - Get table details (schema, columns, properties): - invoke_databricks_cli 'tables get ..
' - - -5. EXPLORING WORKSPACE FILES - List workspace files and notebooks: - invoke_databricks_cli 'workspace list ' - - Export a notebook: - invoke_databricks_cli 'workspace export ' - - -Getting Started: -- Use the commands above to explore what resources exist in the workspace -- All commands support --output json for programmatic access -- Remember to add --profile when working with non-default workspaces - - -WORKFLOW PATTERNS FOR DATABRICKS PROJECTS -========================================== - -Creating a New Databricks Project: - When to use: Building a new project from scratch, setting up deployment to multiple environments - Tools sequence: - 1. init_project (creates proper project structure with templates) - 2. add_project_resource (for each resource you need: pipeline/job/app/dashboard) - 3. analyze_project (provides deployment commands) - 4. invoke_databricks_cli 'bundle validate' - 💡 Tip: Use init_project even if you know YAML syntax - it uses templates and best practices - -Working with Existing Databricks Project: - When to use: databricks.yml file already exists in the directory - Tools sequence: - 1. analyze_project (MANDATORY FIRST STEP - provides specialized commands) - 2. [Make your changes to project files] - 3. invoke_databricks_cli 'bundle validate' - 💡 Tip: ALWAYS call analyze_project before making changes - Databricks projects - require specialized commands that differ from standard Python/Node.js workflows - -Adding Resources to Existing Project: - When to use: Adding pipelines, jobs, apps, or dashboards to an existing project - Tools sequence: - 1. add_project_resource (with type: 'pipeline', 'job', 'app', or 'dashboard') - 2. analyze_project (to get updated deployment commands) - 3. invoke_databricks_cli 'bundle validate' - - -PATTERN MATCHING: If Your Task Mentions... -=========================================== - -"new project" / "create a project" / "Databricks project" / "project structure" - → Use init_project first (don't create files manually!) - → Then add_project_resource for each resource (pipeline/job/app/dashboard) - -"SQL pipeline" / "data pipeline" / "materialized views" / "ETL" / "DLT" - → Use add_project_resource with type='pipeline' or type='job' - -"Databricks app" / "application" / "build an app" - → Use add_project_resource with type='app' - -"dashboard" / "Lakeview dashboard" / "visualization" - → Use add_project_resource with type='dashboard' - -"Databricks job" / "scheduled job" / "workflow" - → Use add_project_resource with type='job' - -"deploy to dev and prod" / "multiple environments" / "dev/staging/prod" - → Use init_project (sets up multi-environment structure automatically) - -"databricks.yml" / "bundle configuration" / "Asset Bundle" - → If creating new: use init_project (don't create manually!) - → If exists already: use analyze_project FIRST before making changes - - -ANTI-PATTERNS TO AVOID -======================= - -❌ DON'T manually create databricks.yml files - ✅ DO use init_project instead - -❌ DON'T run bundle commands without calling analyze_project first - ✅ DO call analyze_project to get the correct specialized commands - -❌ DON'T use regular Bash to run databricks CLI commands - ✅ DO use invoke_databricks_cli (better for user allowlisting) - -❌ DON'T skip explore when planning Databricks work - ✅ DO call explore during planning to get workflow recommendations diff --git a/experimental/aitools/tools/prompts/init_project.tmpl b/experimental/aitools/tools/prompts/init_project.tmpl deleted file mode 100644 index 1e04ed2168..0000000000 --- a/experimental/aitools/tools/prompts/init_project.tmpl +++ /dev/null @@ -1,14 +0,0 @@ -{{- /* - * Success message shown after initializing a new Databricks project. - * - * This template is used when init_project creates a new empty project. - * - */ -}} - -Project '{{.ProjectName}}' initialized successfully at: {{.ProjectPath}} - -⚠️ IMPORTANT: This is an EMPTY project with NO resources (no apps, jobs, pipelines, or dashboards)! - -If the user asked you to create a specific resource (like "create an app" or "create a job"), you MUST now call the add_project_resource tool to add it! - -Example: add_project_resource(project_path="{{.ProjectPath}}", type="app", name="my_app", template="nodejs-fastapi-hello-world-app") diff --git a/experimental/aitools/tools/prompts/jobs.tmpl b/experimental/aitools/tools/prompts/jobs.tmpl deleted file mode 100644 index 82dfa8fa92..0000000000 --- a/experimental/aitools/tools/prompts/jobs.tmpl +++ /dev/null @@ -1,17 +0,0 @@ -{{- /* - * Placeholder instructions for working with jobs. - * - * This template is only appended to analyze_project.tmpl output when jobs - * (resources/*.job.yml files) are detected in the project. It provides - * job-specific guidance for running and deploying jobs. - * - * TODO: these instructions should be refined and benchmarked - * - */ -}} - -Working with Jobs ------------------ -- Jobs are scheduled workflows that can run Python notebooks or SQL queries -- Job source code is in src// for Python jobs or src/*.sql for SQL jobs -- Run jobs using: invoke_databricks_cli(command="bundle run ", working_directory="") -- Deploy jobs using: invoke_databricks_cli(command="bundle deploy --target dev", working_directory="") diff --git a/experimental/aitools/tools/prompts/pipelines.tmpl b/experimental/aitools/tools/prompts/pipelines.tmpl deleted file mode 100644 index 746217b794..0000000000 --- a/experimental/aitools/tools/prompts/pipelines.tmpl +++ /dev/null @@ -1,17 +0,0 @@ -{{- /* - * Placeholder instructions for working with Spark Declarative Pipelines. - * - * This template is only appended to analyze_project.tmpl output when pipelines - * (resources/*.pipeline.yml files) are detected in the project. It provides - * pipeline-specific guidance for validating and deploying Lakeflow pipelines. - * - * TODO: these instructions should be refined and benchmarked - * - */ -}} - -Working with Pipelines ----------------------- -- Pipelines are Lakeflow Spark Declarative pipelines (data processing workflows) -- Pipeline source code is in src// -- Validate pipeline definitions before deployment using: invoke_databricks_cli(command="bundle run --validate-only", working_directory="") -- Deploy pipelines using: invoke_databricks_cli(command="bundle deploy --target dev", working_directory="") diff --git a/experimental/aitools/tools/prompts/prompts.go b/experimental/aitools/tools/prompts/prompts.go deleted file mode 100644 index 44ab5b7c0a..0000000000 --- a/experimental/aitools/tools/prompts/prompts.go +++ /dev/null @@ -1,60 +0,0 @@ -package prompts - -import ( - "bytes" - "embed" - "fmt" - "text/template" -) - -//go:embed *.tmpl -var promptTemplates embed.FS - -// ExecuteTemplate loads and executes a template with the given name and data. -func ExecuteTemplate(name string, data any) (string, error) { - tmplContent, err := promptTemplates.ReadFile(name) - if err != nil { - return "", fmt.Errorf("failed to read template %s: %w", name, err) - } - - tmpl, err := template.New(name).Parse(string(tmplContent)) - if err != nil { - return "", fmt.Errorf("failed to parse template %s: %w", name, err) - } - - var buf bytes.Buffer - if err := tmpl.Execute(&buf, data); err != nil { - return "", fmt.Errorf("failed to execute template %s: %w", name, err) - } - - return buf.String(), nil -} - -// MustExecuteTemplate is like ExecuteTemplate but panics on error. -// Use this only when template execution errors are programming errors. -func MustExecuteTemplate(name string, data any) string { - result, err := ExecuteTemplate(name, data) - if err != nil { - panic(err) - } - return result -} - -// LoadTemplate loads a template without executing it. -// Returns the raw template content as a string. -func LoadTemplate(name string) (string, error) { - tmplContent, err := promptTemplates.ReadFile(name) - if err != nil { - return "", fmt.Errorf("failed to read template %s: %w", name, err) - } - return string(tmplContent), nil -} - -// MustLoadTemplate is like LoadTemplate but panics on error. -func MustLoadTemplate(name string) string { - result, err := LoadTemplate(name) - if err != nil { - panic(err) - } - return result -} diff --git a/experimental/aitools/tools/resources/apps.go b/experimental/aitools/tools/resources/apps.go deleted file mode 100644 index ec54cea7ba..0000000000 --- a/experimental/aitools/tools/resources/apps.go +++ /dev/null @@ -1,95 +0,0 @@ -package resources - -import ( - "context" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/databricks/cli/experimental/aitools/tools/prompts" -) - -type appHandler struct{} - -func (h *appHandler) AddToProject(ctx context.Context, args AddProjectResourceArgs) (string, error) { - // FIXME: This should rely on the databricks bundle generate command to handle all template scaffolding. - - tmpDir, cleanup, err := CloneTemplateRepo(ctx, "https://github.com/databricks/app-templates") - if err != nil { - return "", err - } - defer cleanup() - - if args.Template == "" { - // List available templates - entries, err := os.ReadDir(tmpDir) - if err != nil { - return "", fmt.Errorf("failed to read app-templates directory: %w", err) - } - - var templateList []string - for _, entry := range entries { - if entry.IsDir() && !strings.HasPrefix(entry.Name(), ".") { - templateList = append(templateList, entry.Name()) - } - } - - templateListStr := formatTemplateList(templateList) - errorMsg := prompts.MustExecuteTemplate("apps_pick_a_template.tmpl", map[string]string{ - "TemplateList": templateListStr, - "ProjectPath": args.ProjectPath, - "Name": args.Name, - }) - return "", fmt.Errorf("%s", errorMsg) - } - - templateSrc := filepath.Join(tmpDir, args.Template) - if _, err := os.Stat(templateSrc); os.IsNotExist(err) { - return "", fmt.Errorf("template '%s' not found in app-templates repository", args.Template) - } - - appDest := filepath.Join(args.ProjectPath, args.Name) - if err := CopyDir(templateSrc, appDest); err != nil { - return "", fmt.Errorf("failed to copy app template: %w", err) - } - - replacements := map[string]string{ - args.Template: args.Name, - } - if err := ReplaceInDirectory(appDest, replacements); err != nil { - return "", fmt.Errorf("failed to replace template references: %w", err) - } - - resourceYAML := fmt.Sprintf(`resources: - apps: - %s: - name: %s - description: Databricks app created from %s template - source_code_path: ../%s -`, args.Name, args.Name, args.Template, args.Name) - - resourceFile := filepath.Join(args.ProjectPath, "resources", args.Name+".app.yml") - if err := os.WriteFile(resourceFile, []byte(resourceYAML), 0o644); err != nil { - return "", fmt.Errorf("failed to create resource file: %w", err) - } - - return "", nil -} - -func (h *appHandler) GetGuidancePrompt(projectPath, warehouseID, warehouseName string) string { - return prompts.MustExecuteTemplate("apps.tmpl", map[string]string{ - "WarehouseID": warehouseID, - "WarehouseName": warehouseName, - }) -} - -func formatTemplateList(templates []string) string { - var result strings.Builder - for _, t := range templates { - result.WriteString("- ") - result.WriteString(t) - result.WriteString("\n") - } - return result.String() -} diff --git a/experimental/aitools/tools/resources/common.go b/experimental/aitools/tools/resources/common.go deleted file mode 100644 index acd88a8740..0000000000 --- a/experimental/aitools/tools/resources/common.go +++ /dev/null @@ -1,169 +0,0 @@ -package resources - -import ( - "context" - "fmt" - "os" - "os/exec" - "path/filepath" - "strings" -) - -// AddProjectResourceArgs represents the arguments for adding a resource. -type AddProjectResourceArgs struct { - ProjectPath string `json:"project_path"` - Type string `json:"type"` - Name string `json:"name"` - Template string `json:"template,omitempty"` -} - -// ResourceHandler defines the interface for resource-specific operations. -type ResourceHandler interface { - // AddToProject adds the resource to a project. - AddToProject(ctx context.Context, args AddProjectResourceArgs) (string, error) - // GetGuidancePrompt returns resource-specific guidance for the AI agent. - // warehouseID and warehouseName are optional (empty strings if not available). - GetGuidancePrompt(projectPath, warehouseID, warehouseName string) string -} - -// CloneTemplateRepo clones a GitHub repository to a temporary directory. -func CloneTemplateRepo(ctx context.Context, repoURL string) (string, func(), error) { - tmpDir, err := os.MkdirTemp("", filepath.Base(repoURL)+"-*") - if err != nil { - return "", nil, fmt.Errorf("failed to create temp directory: %w", err) - } - - cleanup := func() { os.RemoveAll(tmpDir) } - - cmd := exec.CommandContext(ctx, "git", "clone", "--depth", "1", repoURL, tmpDir) - if output, err := cmd.CombinedOutput(); err != nil { - cleanup() - return "", nil, fmt.Errorf("failed to clone %s: %w\nOutput: %s", filepath.Base(repoURL), err, string(output)) - } - - return tmpDir, cleanup, nil -} - -// CopyResourceFile copies and renames a resource YAML file from template to project. -func CopyResourceFile(resourceSrc, projectPath, resourceName, suffix string, replacements map[string]string) error { - files, err := os.ReadDir(resourceSrc) - if err != nil { - return fmt.Errorf("failed to read template resources: %w", err) - } - - for _, file := range files { - if strings.HasSuffix(file.Name(), suffix) { - content, err := os.ReadFile(filepath.Join(resourceSrc, file.Name())) - if err != nil { - return fmt.Errorf("failed to read template: %w", err) - } - - modified := string(content) - for old, new := range replacements { - modified = strings.ReplaceAll(modified, old, new) - } - - dstFile := filepath.Join(projectPath, "resources", resourceName+suffix) - return os.WriteFile(dstFile, []byte(modified), 0o644) - } - } - return fmt.Errorf("no %s file found in template", suffix) -} - -// ReplaceInDirectory walks through a directory and replaces strings in all text files. -func ReplaceInDirectory(dir string, replacements map[string]string) error { - return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error { - if err != nil { - return err - } - - if info.IsDir() { - return nil - } - - ext := strings.ToLower(filepath.Ext(path)) - textExts := []string{".py", ".sql", ".yml", ".yaml", ".md", ".txt", ".toml", ".json", ".sh"} - isTextFile := false - for _, textExt := range textExts { - if ext == textExt { - isTextFile = true - break - } - } - - if !isTextFile { - return nil - } - - content, err := os.ReadFile(path) - if err != nil { - return err - } - - modified := string(content) - for old, new := range replacements { - modified = strings.ReplaceAll(modified, old, new) - } - - if modified != string(content) { - return os.WriteFile(path, []byte(modified), 0o644) - } - return nil - }) -} - -// CopyDir recursively copies a directory from src to dst. -func CopyDir(src, dst string) error { - return filepath.Walk(src, func(path string, info os.FileInfo, err error) error { - if err != nil { - return err - } - - relPath, err := filepath.Rel(src, path) - if err != nil { - return err - } - - dstPath := filepath.Join(dst, relPath) - - if info.IsDir() { - return os.MkdirAll(dstPath, info.Mode()) - } - - content, err := os.ReadFile(path) - if err != nil { - return err - } - - return os.WriteFile(dstPath, content, info.Mode()) - }) -} - -// ValidateLanguageTemplate validates that a template is specified and is either 'python' or 'sql'. -func ValidateLanguageTemplate(template, resourceType string) error { - if template == "" { - return fmt.Errorf("template is required for %ss\n\nPlease specify the language: 'python' or 'sql'", resourceType) - } - - if template != "python" && template != "sql" { - return fmt.Errorf("invalid template for %s: %s. Must be 'python' or 'sql'", resourceType, template) - } - - return nil -} - -// GetResourceHandler returns the ResourceHandler for the given resource type. -func GetResourceHandler(resourceType string) ResourceHandler { - switch resourceType { - case "app": - return &appHandler{} - case "job": - return &jobHandler{} - case "pipeline": - return &pipelineHandler{} - case "dashboard": - return &dashboardHandler{} - default: - return nil - } -} diff --git a/experimental/aitools/tools/resources/common_test.go b/experimental/aitools/tools/resources/common_test.go deleted file mode 100644 index 955bcd1934..0000000000 --- a/experimental/aitools/tools/resources/common_test.go +++ /dev/null @@ -1,54 +0,0 @@ -package resources - -import ( - "os" - "path/filepath" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestCopyResourceFile(t *testing.T) { - // Create temp directories for source and destination - tmpDir := t.TempDir() - srcDir := filepath.Join(tmpDir, "src", "resources") - dstDir := filepath.Join(tmpDir, "dst") - resourcesDir := filepath.Join(dstDir, "resources") - - require.NoError(t, os.MkdirAll(srcDir, 0o755)) - require.NoError(t, os.MkdirAll(resourcesDir, 0o755)) - - // Create a test resource file - testContent := "name: test_job\npackage: test_package" - testFile := filepath.Join(srcDir, "template.job.yml") - require.NoError(t, os.WriteFile(testFile, []byte(testContent), 0o644)) - - // Test copying with replacements - replacements := map[string]string{ - "test_job": "my_job", - "test_package": "my_package", - } - err := CopyResourceFile(srcDir, dstDir, "my_job", ".job.yml", replacements) - assert.NoError(t, err) - - // Verify the file was copied and modified - dstFile := filepath.Join(resourcesDir, "my_job.job.yml") - content, err := os.ReadFile(dstFile) - require.NoError(t, err) - assert.Contains(t, string(content), "my_job") - assert.Contains(t, string(content), "my_package") - assert.NotContains(t, string(content), "test_job") - assert.NotContains(t, string(content), "test_package") -} - -func TestCopyResourceFile_NotFound(t *testing.T) { - tmpDir := t.TempDir() - srcDir := filepath.Join(tmpDir, "src", "resources") - require.NoError(t, os.MkdirAll(srcDir, 0o755)) - - // Test with non-existent file type - err := CopyResourceFile(srcDir, tmpDir, "test", ".nonexistent.yml", nil) - assert.Error(t, err) - assert.Contains(t, err.Error(), "no .nonexistent.yml file found") -} diff --git a/experimental/aitools/tools/resources/dashboards.go b/experimental/aitools/tools/resources/dashboards.go deleted file mode 100644 index 084679bd74..0000000000 --- a/experimental/aitools/tools/resources/dashboards.go +++ /dev/null @@ -1,79 +0,0 @@ -package resources - -import ( - "context" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/databricks/cli/experimental/aitools/tools/prompts" -) - -type dashboardHandler struct{} - -func (h *dashboardHandler) AddToProject(ctx context.Context, args AddProjectResourceArgs) (string, error) { - // FIXME: This should rely on the databricks bundle generate command to handle all template scaffolding. - - tmpDir, cleanup, err := CloneTemplateRepo(ctx, "https://github.com/databricks/bundle-examples") - if err != nil { - return "", err - } - defer cleanup() - - templateSrc := filepath.Join(tmpDir, "knowledge_base", "dashboard_nyc_taxi") - - // Dashboard templates use the file name (without extension) as the resource name - resourceFiles, err := os.ReadDir(filepath.Join(templateSrc, "resources")) - if err != nil { - return "", fmt.Errorf("failed to read template resources: %w", err) - } - - var oldName string - for _, file := range resourceFiles { - if strings.HasSuffix(file.Name(), ".dashboard.yml") { - oldName = strings.TrimSuffix(file.Name(), ".dashboard.yml") - break - } - } - - replacements := map[string]string{oldName: args.Name} - if err := CopyResourceFile(filepath.Join(templateSrc, "resources"), args.ProjectPath, args.Name, ".dashboard.yml", replacements); err != nil { - return "", err - } - - srcDir := filepath.Join(args.ProjectPath, "src") - if err := os.MkdirAll(srcDir, 0o755); err != nil { - return "", fmt.Errorf("failed to create src directory: %w", err) - } - - dashFiles, err := filepath.Glob(filepath.Join(templateSrc, "src", "*.lvdash.json")) - if err != nil { - return "", fmt.Errorf("failed to find dashboard files: %w", err) - } - - for _, dashFile := range dashFiles { - basename := filepath.Base(dashFile) - // Extract the old name from the filename (e.g., "nyc_taxi.lvdash.json" -> "nyc_taxi") - oldName := strings.TrimSuffix(basename, ".lvdash.json") - newName := args.Name + ".lvdash.json" - dstFile := filepath.Join(srcDir, newName) - - content, err := os.ReadFile(dashFile) - if err != nil { - return "", fmt.Errorf("failed to read dashboard JSON: %w", err) - } - - modified := strings.ReplaceAll(string(content), oldName, args.Name) - - if err := os.WriteFile(dstFile, []byte(modified), 0o644); err != nil { - return "", fmt.Errorf("failed to write dashboard JSON: %w", err) - } - } - - return "", nil -} - -func (h *dashboardHandler) GetGuidancePrompt(projectPath, warehouseID, warehouseName string) string { - return prompts.MustLoadTemplate("dashboards.tmpl") -} diff --git a/experimental/aitools/tools/resources/jobs.go b/experimental/aitools/tools/resources/jobs.go deleted file mode 100644 index bb1f94b2a1..0000000000 --- a/experimental/aitools/tools/resources/jobs.go +++ /dev/null @@ -1,123 +0,0 @@ -package resources - -import ( - "context" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/databricks/cli/experimental/aitools/tools/prompts" -) - -type jobHandler struct{} - -func (h *jobHandler) AddToProject(ctx context.Context, args AddProjectResourceArgs) (string, error) { - // FIXME: This should rely on the databricks bundle generate command to handle all template scaffolding. - - if err := ValidateLanguageTemplate(args.Template, "job"); err != nil { - return "", err - } - - if args.Template == "python" { - return "", addPythonJob(ctx, args) - } - return "", addSQLJob(ctx, args) -} - -func (h *jobHandler) GetGuidancePrompt(projectPath, warehouseID, warehouseName string) string { - return prompts.MustLoadTemplate("jobs.tmpl") -} - -func addPythonJob(ctx context.Context, args AddProjectResourceArgs) error { - tmpDir, cleanup, err := CloneTemplateRepo(ctx, "https://github.com/databricks/bundle-examples") - if err != nil { - return err - } - defer cleanup() - - templateName := "default_python" - templateSrc := filepath.Join(tmpDir, templateName) - replacements := map[string]string{templateName: args.Name} - - if err := CopyResourceFile(filepath.Join(templateSrc, "resources"), args.ProjectPath, args.Name, ".job.yml", replacements); err != nil { - return err - } - - srcDir := filepath.Join(args.ProjectPath, "src") - if err := os.MkdirAll(srcDir, 0o755); err != nil { - return fmt.Errorf("failed to create src directory: %w", err) - } - - pythonSrc := filepath.Join(templateSrc, "src", templateName) - pythonDest := filepath.Join(srcDir, args.Name) - if err := CopyDir(pythonSrc, pythonDest); err != nil { - return fmt.Errorf("failed to copy python source: %w", err) - } - - if err := ReplaceInDirectory(pythonDest, replacements); err != nil { - return fmt.Errorf("failed to replace template references in source: %w", err) - } - - testsDest := filepath.Join(args.ProjectPath, "tests") - if _, err := os.Stat(testsDest); os.IsNotExist(err) { - testsSrc := filepath.Join(templateSrc, "tests") - if err := CopyDir(testsSrc, testsDest); err != nil { - fmt.Fprintf(os.Stderr, "Warning: failed to copy tests: %v\n", err) - } else if err := ReplaceInDirectory(testsDest, replacements); err != nil { - fmt.Fprintf(os.Stderr, "Warning: failed to replace template references in tests: %v\n", err) - } - } - return nil -} - -func addSQLJob(ctx context.Context, args AddProjectResourceArgs) error { - tmpDir, cleanup, err := CloneTemplateRepo(ctx, "https://github.com/databricks/bundle-examples") - if err != nil { - return err - } - defer cleanup() - - templateName := "default_sql" - templateSrc := filepath.Join(tmpDir, templateName) - replacements := map[string]string{templateName: args.Name} - - if err := CopyResourceFile(filepath.Join(templateSrc, "resources"), args.ProjectPath, args.Name, ".job.yml", replacements); err != nil { - return err - } - - srcDir := filepath.Join(args.ProjectPath, "src") - if err := os.MkdirAll(srcDir, 0o755); err != nil { - return fmt.Errorf("failed to create src directory: %w", err) - } - - sqlFiles, err := filepath.Glob(filepath.Join(templateSrc, "src", "*.sql")) - if err != nil { - return fmt.Errorf("failed to find SQL files: %w", err) - } - - for _, sqlFile := range sqlFiles { - basename := filepath.Base(sqlFile) - newName := strings.ReplaceAll(basename, templateName, args.Name) - dstFile := filepath.Join(srcDir, newName) - - if _, err := os.Stat(dstFile); err == nil { - continue - } - - content, err := os.ReadFile(sqlFile) - if err != nil { - return fmt.Errorf("failed to read SQL file: %w", err) - } - - modified := string(content) - for old, new := range replacements { - modified = strings.ReplaceAll(modified, old, new) - } - - if err := os.WriteFile(dstFile, []byte(modified), 0o644); err != nil { - return fmt.Errorf("failed to write SQL file: %w", err) - } - } - return nil -} diff --git a/experimental/aitools/tools/resources/pipelines.go b/experimental/aitools/tools/resources/pipelines.go deleted file mode 100644 index 934627d2a8..0000000000 --- a/experimental/aitools/tools/resources/pipelines.go +++ /dev/null @@ -1,72 +0,0 @@ -package resources - -import ( - "context" - "fmt" - "os" - "path/filepath" - - "github.com/databricks/cli/experimental/aitools/tools/prompts" -) - -type pipelineHandler struct{} - -func (h *pipelineHandler) AddToProject(ctx context.Context, args AddProjectResourceArgs) (string, error) { - // FIXME: This should rely on the databricks bundle generate command to handle all template scaffolding. - - if err := ValidateLanguageTemplate(args.Template, "pipeline"); err != nil { - return "", err - } - - tmpDir, cleanup, err := CloneTemplateRepo(ctx, "https://github.com/databricks/bundle-examples") - if err != nil { - return "", err - } - defer cleanup() - - templateName := "lakeflow_pipelines_" + args.Template - templateSrc := filepath.Join(tmpDir, templateName) - - // Copy source files first - use the actual directory name from bundle-examples - srcDir := filepath.Join(templateSrc, "src") - srcEntries, err := os.ReadDir(srcDir) - if err != nil { - return "", fmt.Errorf("failed to read pipeline src directory: %w", err) - } - - // Find the first directory in src/ - var srcSubdir string - for _, entry := range srcEntries { - if entry.IsDir() { - srcSubdir = entry.Name() - break - } - } - if srcSubdir == "" { - return "", fmt.Errorf("no source directory found in %s", srcDir) - } - - srcPattern := filepath.Join(srcDir, srcSubdir) - srcDest := filepath.Join(args.ProjectPath, "src", args.Name) - if err := CopyDir(srcPattern, srcDest); err != nil { - return "", fmt.Errorf("failed to copy pipeline source: %w", err) - } - - replacements := map[string]string{ - templateName: args.Name, - srcSubdir: args.Name, - } - if err := CopyResourceFile(filepath.Join(templateSrc, "resources"), args.ProjectPath, args.Name, ".pipeline.yml", replacements); err != nil { - return "", err - } - - if err := ReplaceInDirectory(srcDest, replacements); err != nil { - return "", fmt.Errorf("failed to replace template references: %w", err) - } - - return "", nil -} - -func (h *pipelineHandler) GetGuidancePrompt(projectPath, warehouseID, warehouseName string) string { - return prompts.MustLoadTemplate("pipelines.tmpl") -} diff --git a/experimental/aitools/uninstall.go b/experimental/aitools/uninstall.go deleted file mode 100644 index fed4baf8ae..0000000000 --- a/experimental/aitools/uninstall.go +++ /dev/null @@ -1,27 +0,0 @@ -package aitools - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -func newUninstallCmd() *cobra.Command { - return &cobra.Command{ - Use: "uninstall", - Short: "Show instructions for uninstalling the MCP server", - Long: `Show instructions for uninstalling the Databricks CLI MCP server from coding agents.`, - Run: func(cmd *cobra.Command, args []string) { - // Currently we just provide instructions for users to uninstall manually via their agent. - fmt.Print(` -To uninstall the Databricks CLI MCP server, please ask your coding agent to remove it. - -For Claude Code, you can also use: - claude mcp remove databricks-aitools - -For Cursor, you can also manually remove the entry from: - ~/.cursor/mcp.json -`) - }, - } -} diff --git a/tools/testmask/targets.go b/tools/testmask/targets.go index 0ab4f71101..404b07e82c 100644 --- a/tools/testmask/targets.go +++ b/tools/testmask/targets.go @@ -24,12 +24,6 @@ var fileTargetMappings = []targetMapping{ }), target: "test", }, - { - prefixes: slices.Concat(goTriggerPatterns, []string{ - "experimental/aitools/", - }), - target: "test-exp-aitools", - }, { prefixes: slices.Concat(goTriggerPatterns, []string{ "experimental/apps-mcp/", diff --git a/tools/testmask/targets_test.go b/tools/testmask/targets_test.go index 8b693aeb53..beace5779e 100644 --- a/tools/testmask/targets_test.go +++ b/tools/testmask/targets_test.go @@ -24,13 +24,6 @@ func TestGetTargets(t *testing.T) { }, targets: []string{"test-exp-ssh"}, }, - { - name: "experimental_aitools", - files: []string{ - "experimental/aitools/server.go", - }, - targets: []string{"test-exp-aitools"}, - }, { name: "pipelines", files: []string{ @@ -51,14 +44,14 @@ func TestGetTargets(t *testing.T) { files: []string{ "go.mod", }, - targets: []string{"test", "test-exp-aitools", "test-exp-apps-mcp", "test-exp-ssh", "test-pipelines"}, + targets: []string{"test", "test-exp-apps-mcp", "test-exp-ssh", "test-pipelines"}, }, { name: "go_sum_triggers_all", files: []string{ "go.sum", }, - targets: []string{"test", "test-exp-aitools", "test-exp-apps-mcp", "test-exp-ssh", "test-pipelines"}, + targets: []string{"test", "test-exp-apps-mcp", "test-exp-ssh", "test-pipelines"}, }, { name: "go_mod_with_other_files_triggers_all", @@ -66,7 +59,7 @@ func TestGetTargets(t *testing.T) { "experimental/ssh/main.go", "go.mod", }, - targets: []string{"test", "test-exp-aitools", "test-exp-apps-mcp", "test-exp-ssh", "test-pipelines"}, + targets: []string{"test", "test-exp-apps-mcp", "test-exp-ssh", "test-pipelines"}, }, { name: "empty_files",