Consolidate compute tools: 19 → 4 MCP tools for execution & compute management (#278)

* Add serverless code execution via Jobs API (no cluster required)

Adds a new `run_code_on_serverless()` function that executes Python or SQL
code on Databricks serverless compute using the Jobs API `runs/submit`
endpoint. No interactive cluster is required.

The implementation:
- Uploads code as a temporary notebook to the workspace
- Submits a one-time run with serverless compute (environments + environment_key pattern)
- Waits for completion and retrieves output via get_run_output
- Cleans up temporary workspace files after execution
- Returns a typed ServerlessRunResult with output, error, run_id, run_url, and timing

New files and changes:
- databricks-tools-core: compute/serverless.py (core module)
- databricks-tools-core: compute/__init__.py (exports)
- databricks-mcp-server: tools/compute.py (MCP tool wrapper)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Fix serverless code runner: capture error tracebacks, add skill file, improve docs

- Retrieve actual Python traceback on failure instead of generic "Workload
  failed" message by fetching run output in the exception handler
- Fix Optional[str] type annotation for run_name in MCP wrapper
- Document SQL SELECT output limitation in all docstrings
- Reframe tool as Python-first; clarify SQL is niche (DDL/DML only, use
  execute_sql for queries — works with serverless SQL warehouses)
- Add databricks-serverless-compute skill file with decision matrix,
  output capture behavior, limitations, and examples

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Add .ipynb (Jupyter notebook) support to serverless code runner

Auto-detects .ipynb JSON content and uploads via Databricks native
Jupyter import (ImportFormat.JUPYTER), enabling users to run local
Jupyter notebooks on serverless compute without conversion.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* Consolidate into databricks-execution-compute skill with multi-language and workspace_path support

- Rename databricks-serverless-compute → databricks-execution-compute covering all three
  execution tools (execute_databricks_command, run_file_on_databricks, run_code_on_serverless)
- Rename run_python_file_on_databricks → run_file_on_databricks with language auto-detection
  from file extension (.py, .scala, .sql, .r); old name kept as alias
- Add workspace_path param to run_file_on_databricks and run_code_on_serverless for
  persistent mode (saves notebook to workspace) vs ephemeral (default, temp cleanup)
- Add comprehensive integration tests (34 tests) covering classic cluster execution,
  serverless Python/SQL, ephemeral vs persistent modes, input validation, and error handling
- Update MCP tool layer with new params and empty-string-to-None coercion
- Update install_skills.sh with new skill name and description

Co-authored-by: Isaac

* Add databricks-manage-compute skill for cluster and warehouse lifecycle management

Core functions (manage.py): create/modify/terminate/delete clusters with opinionated
defaults (auto-pick LTS DBR, reasonable node type, SINGLE_USER mode, 120min auto-term),
plus create/modify/delete SQL warehouses. List helpers for node types and spark versions.
MCP tool wrappers with destructive-action warnings in docstrings. SKILL.md with decision
matrix, tool reference tables, and examples. Integration tests validated against e2-demo-field-eng.

Co-authored-by: Isaac

* Consolidate 19 compute MCP tools into 4 per PR feedback

Address Cal's feedback to reduce tool count and LLM parsing bloat.
Uses action/resource params to determine behavior (like SDP's
create_and_update pattern).

New tools:
- execute_code (replaces 3 execution tools)
- manage_cluster (replaces 5 cluster tools)
- manage_sql_warehouse (replaces 3 warehouse tools)
- list_compute (replaces 5 listing/inspection tools)

Also removes the separate databricks-manage-compute skill (merged into
databricks-execution-compute) and adds 41 unit tests for routing logic.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: George Theodosopoulos <george.theodosopoulos@databricks.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Cal Reynolds <49540501+calreynolds@users.noreply.github.com>

Author: 4 people

databricks-mcp-server/databricks_mcp_server/tools/compute.py

@@ -0,0 +1,392 @@
+"""
+Unit tests for consolidated compute tools.
+
+Tests the MCP tool wrapper routing logic without hitting Databricks APIs.
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock
+from databricks_mcp_server.tools.compute import (
+    execute_code,
+    manage_cluster,
+    manage_sql_warehouse,
+    list_compute,
+)
+
+
+# ---------------------------------------------------------------------------
+# execute_code routing tests
+# ---------------------------------------------------------------------------
+
+
+class TestExecuteCodeRouting:
+    """Test that execute_code routes to the correct backend."""
+
+    def test_requires_code_or_file_path(self):
+        result = execute_code()
+        assert result["success"] is False
+        assert "code" in result["error"].lower() or "file_path" in result["error"].lower()
+
+    def test_empty_strings_treated_as_none(self):
+        result = execute_code(code="", file_path="")
+        assert result["success"] is False
+        assert "code" in result["error"].lower() or "file_path" in result["error"].lower()
+
+    @patch("databricks_mcp_server.tools.compute._run_code_on_serverless")
+    def test_auto_routes_to_serverless_for_python(self, mock_serverless):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True, "output": "hello"}
+        mock_serverless.return_value = mock_result
+
+        execute_code(code="print('hi')", compute_type="auto")
+
+        mock_serverless.assert_called_once()
+        call_kwargs = mock_serverless.call_args[1]
+        assert call_kwargs["code"] == "print('hi')"
+        assert call_kwargs["language"] == "python"
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_auto_routes_to_cluster_with_cluster_id(self, mock_cluster):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_cluster.return_value = mock_result
+
+        execute_code(code="print('hi')", cluster_id="abc-123")
+
+        mock_cluster.assert_called_once()
+        assert mock_cluster.call_args[1]["cluster_id"] == "abc-123"
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_auto_routes_to_cluster_with_context_id(self, mock_cluster):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_cluster.return_value = mock_result
+
+        execute_code(code="print('hi')", context_id="ctx-456")
+
+        mock_cluster.assert_called_once()
+        assert mock_cluster.call_args[1]["context_id"] == "ctx-456"
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_auto_routes_to_cluster_for_scala(self, mock_cluster):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_cluster.return_value = mock_result
+
+        execute_code(code="println(42)", language="scala")
+
+        mock_cluster.assert_called_once()
+        assert mock_cluster.call_args[1]["language"] == "scala"
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_auto_routes_to_cluster_for_r(self, mock_cluster):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_cluster.return_value = mock_result
+
+        execute_code(code="print(42)", language="r")
+
+        mock_cluster.assert_called_once()
+
+    @patch("databricks_mcp_server.tools.compute._run_code_on_serverless")
+    def test_explicit_serverless(self, mock_serverless):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_serverless.return_value = mock_result
+
+        execute_code(code="print('hi')", compute_type="serverless")
+
+        mock_serverless.assert_called_once()
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_explicit_cluster(self, mock_cluster):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_cluster.return_value = mock_result
+
+        execute_code(code="print('hi')", compute_type="cluster")
+
+        mock_cluster.assert_called_once()
+
+    @patch("databricks_mcp_server.tools.compute._run_code_on_serverless")
+    def test_serverless_default_timeout(self, mock_serverless):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_serverless.return_value = mock_result
+
+        execute_code(code="x", compute_type="serverless")
+
+        assert mock_serverless.call_args[1]["timeout"] == 1800
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_cluster_default_timeout(self, mock_cluster):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_cluster.return_value = mock_result
+
+        execute_code(code="x", compute_type="cluster")
+
+        assert mock_cluster.call_args[1]["timeout"] == 120
+
+    @patch("databricks_mcp_server.tools.compute._run_code_on_serverless")
+    def test_workspace_path_passed_to_serverless(self, mock_serverless):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_serverless.return_value = mock_result
+
+        execute_code(code="x", compute_type="serverless", workspace_path="/Workspace/Users/a/b")
+
+        call_kwargs = mock_serverless.call_args[1]
+        assert call_kwargs["workspace_path"] == "/Workspace/Users/a/b"
+        assert call_kwargs["cleanup"] is False
+
+    @patch("databricks_mcp_server.tools.compute._run_file_on_databricks")
+    def test_file_path_on_cluster(self, mock_run_file):
+        mock_result = MagicMock()
+        mock_result.to_dict.return_value = {"success": True}
+        mock_run_file.return_value = mock_result
+
+        execute_code(file_path="/tmp/test.py", compute_type="cluster")
+
+        mock_run_file.assert_called_once()
+        assert mock_run_file.call_args[1]["file_path"] == "/tmp/test.py"
+
+    def test_file_path_not_found_serverless(self):
+        result = execute_code(file_path="/nonexistent/file.py", compute_type="serverless")
+        assert result["success"] is False
+        assert "not found" in result["error"].lower()
+
+    @patch("databricks_mcp_server.tools.compute._execute_databricks_command")
+    def test_no_running_cluster_error(self, mock_cluster):
+        from databricks_tools_core.compute import NoRunningClusterError
+        mock_cluster.side_effect = NoRunningClusterError(
+            available_clusters=[],
+            skipped_clusters=[],
+            startable_clusters=[{"cluster_id": "abc", "cluster_name": "test", "state": "TERMINATED"}],
+        )
+
+        result = execute_code(code="x", compute_type="cluster")
+
+        assert result["success"] is False
+        assert "startable_clusters" in result
+        assert len(result["startable_clusters"]) == 1
+
+
+# ---------------------------------------------------------------------------
+# manage_cluster routing tests
+# ---------------------------------------------------------------------------
+
+
+class TestManageCluster:
+    """Test manage_cluster action routing."""
+
+    def test_invalid_action(self):
+        result = manage_cluster(action="explode")
+        assert result["success"] is False
+        assert "unknown action" in result["error"].lower()
+
+    def test_create_requires_name(self):
+        result = manage_cluster(action="create")
+        assert result["success"] is False
+        assert "name" in result["error"].lower()
+
+    def test_modify_requires_cluster_id(self):
+        result = manage_cluster(action="modify")
+        assert result["success"] is False
+        assert "cluster_id" in result["error"].lower()
+
+    def test_start_requires_cluster_id(self):
+        result = manage_cluster(action="start")
+        assert result["success"] is False
+        assert "cluster_id" in result["error"].lower()
+
+    def test_terminate_requires_cluster_id(self):
+        result = manage_cluster(action="terminate")
+        assert result["success"] is False
+        assert "cluster_id" in result["error"].lower()
+
+    def test_delete_requires_cluster_id(self):
+        result = manage_cluster(action="delete")
+        assert result["success"] is False
+        assert "cluster_id" in result["error"].lower()
+
+    @patch("databricks_mcp_server.tools.compute._create_cluster")
+    def test_create_routes_correctly(self, mock_create):
+        mock_create.return_value = {"cluster_id": "abc", "state": "PENDING"}
+
+        result = manage_cluster(action="create", name="test-cluster", num_workers=2)
+
+        mock_create.assert_called_once()
+        assert mock_create.call_args[1]["name"] == "test-cluster"
+        assert mock_create.call_args[1]["num_workers"] == 2
+
+    @patch("databricks_mcp_server.tools.compute._modify_cluster")
+    def test_modify_routes_correctly(self, mock_modify):
+        mock_modify.return_value = {"cluster_id": "abc"}
+
+        manage_cluster(action="modify", cluster_id="abc", num_workers=4)
+
+        mock_modify.assert_called_once()
+        assert mock_modify.call_args[1]["cluster_id"] == "abc"
+        assert mock_modify.call_args[1]["num_workers"] == 4
+
+    @patch("databricks_mcp_server.tools.compute._start_cluster")
+    def test_start_routes_correctly(self, mock_start):
+        mock_start.return_value = {"cluster_id": "abc", "state": "PENDING"}
+
+        manage_cluster(action="start", cluster_id="abc")
+
+        mock_start.assert_called_once_with("abc")
+
+    @patch("databricks_mcp_server.tools.compute._terminate_cluster")
+    def test_terminate_routes_correctly(self, mock_terminate):
+        mock_terminate.return_value = {"cluster_id": "abc", "state": "TERMINATING"}
+
+        manage_cluster(action="terminate", cluster_id="abc")
+
+        mock_terminate.assert_called_once_with("abc")
+
+    @patch("databricks_mcp_server.tools.compute._delete_cluster")
+    def test_delete_routes_correctly(self, mock_delete):
+        mock_delete.return_value = {"cluster_id": "abc", "state": "DELETED"}
+
+        manage_cluster(action="delete", cluster_id="abc")
+
+        mock_delete.assert_called_once_with("abc")
+
+    @patch("databricks_mcp_server.tools.compute._create_cluster")
+    def test_create_defaults(self, mock_create):
+        mock_create.return_value = {"cluster_id": "abc"}
+
+        manage_cluster(action="create", name="test")
+
+        call_kwargs = mock_create.call_args[1]
+        assert call_kwargs["num_workers"] == 1
+        assert call_kwargs["autotermination_minutes"] == 120
+
+    @patch("databricks_mcp_server.tools.compute._create_cluster")
+    def test_create_with_spark_conf_json(self, mock_create):
+        mock_create.return_value = {"cluster_id": "abc"}
+
+        manage_cluster(
+            action="create",
+            name="test",
+            spark_conf='{"spark.sql.shuffle.partitions": "8"}',
+        )
+
+        call_kwargs = mock_create.call_args[1]
+        assert call_kwargs["spark_conf"] == {"spark.sql.shuffle.partitions": "8"}
+
+
+# ---------------------------------------------------------------------------
+# manage_sql_warehouse routing tests
+# ---------------------------------------------------------------------------
+
+
+class TestManageSqlWarehouse:
+    """Test manage_sql_warehouse action routing."""
+
+    def test_invalid_action(self):
+        result = manage_sql_warehouse(action="explode")
+        assert result["success"] is False
+
+    def test_create_requires_name(self):
+        result = manage_sql_warehouse(action="create")
+        assert result["success"] is False
+        assert "name" in result["error"].lower()
+
+    def test_modify_requires_warehouse_id(self):
+        result = manage_sql_warehouse(action="modify")
+        assert result["success"] is False
+        assert "warehouse_id" in result["error"].lower()
+
+    def test_delete_requires_warehouse_id(self):
+        result = manage_sql_warehouse(action="delete")
+        assert result["success"] is False
+        assert "warehouse_id" in result["error"].lower()
+
+    @patch("databricks_mcp_server.tools.compute._create_sql_warehouse")
+    def test_create_routes_correctly(self, mock_create):
+        mock_create.return_value = {"warehouse_id": "abc"}
+
+        manage_sql_warehouse(action="create", name="test-wh", size="Medium")
+
+        mock_create.assert_called_once()
+        assert mock_create.call_args[1]["name"] == "test-wh"
+        assert mock_create.call_args[1]["size"] == "Medium"
+
+    @patch("databricks_mcp_server.tools.compute._modify_sql_warehouse")
+    def test_modify_routes_correctly(self, mock_modify):
+        mock_modify.return_value = {"warehouse_id": "abc"}
+
+        manage_sql_warehouse(action="modify", warehouse_id="abc", size="Large")
+
+        mock_modify.assert_called_once()
+        assert mock_modify.call_args[1]["size"] == "Large"
+
+    @patch("databricks_mcp_server.tools.compute._delete_sql_warehouse")
+    def test_delete_routes_correctly(self, mock_delete):
+        mock_delete.return_value = {"warehouse_id": "abc"}
+
+        manage_sql_warehouse(action="delete", warehouse_id="abc")
+
+        mock_delete.assert_called_once_with("abc")
+
+
+# ---------------------------------------------------------------------------
+# list_compute routing tests
+# ---------------------------------------------------------------------------
+
+
+class TestListCompute:
+    """Test list_compute resource routing."""
+
+    @patch("databricks_mcp_server.tools.compute._list_clusters")
+    def test_default_lists_clusters(self, mock_list):
+        mock_list.return_value = [{"cluster_id": "abc", "cluster_name": "test"}]
+
+        result = list_compute()
+
+        mock_list.assert_called_once()
+        assert "clusters" in result
+
+    @patch("databricks_mcp_server.tools.compute._get_cluster_status")
+    def test_cluster_id_gets_status(self, mock_status):
+        mock_status.return_value = {"cluster_id": "abc", "state": "RUNNING"}
+
+        result = list_compute(cluster_id="abc")
+
+        mock_status.assert_called_once_with("abc")
+        assert result["state"] == "RUNNING"
+
+    @patch("databricks_mcp_server.tools.compute._get_best_cluster")
+    def test_auto_select(self, mock_best):
+        mock_best.return_value = "best-cluster-id"
+
+        result = list_compute(auto_select=True)
+
+        mock_best.assert_called_once()
+        assert result["cluster_id"] == "best-cluster-id"
+
+    @patch("databricks_mcp_server.tools.compute._list_node_types")
+    def test_node_types(self, mock_nodes):
+        mock_nodes.return_value = [{"node_type_id": "i3.xlarge"}]
+
+        result = list_compute(resource="node_types")
+
+        mock_nodes.assert_called_once()
+        assert "node_types" in result
+
+    @patch("databricks_mcp_server.tools.compute._list_spark_versions")
+    def test_spark_versions(self, mock_versions):
+        mock_versions.return_value = [{"key": "15.4.x-scala2.12"}]
+
+        result = list_compute(resource="spark_versions")
+
+        mock_versions.assert_called_once()
+        assert "spark_versions" in result
+
+    def test_invalid_resource(self):
+        result = list_compute(resource="invalid")
+        assert result["success"] is False
+        assert "unknown resource" in result["error"].lower()