Skip to content

Commit fc72eb7

Browse files
committed
Fix install.sh skill dir append bug and add MCP tool references to jobs and UC skills
Consolidates PRs databricks-solutions#272, databricks-solutions#261, and databricks-solutions#259 into a single PR: - install.sh: Fix Claude skill directory using = instead of +=, which overwrites directories from previously processed tools - databricks-jobs: Add MCP Tool Integration section with manage_jobs and manage_job_runs usage examples - databricks-unity-catalog: Expand MCP Tool Integration with 8 governance tool references (manage_uc_objects, grants, tags, storage, connections, security_policies, monitors, sharing) - Fix pre-existing ruff format issues in auth.py and test_sql.py
1 parent dd3684e commit fc72eb7

5 files changed

Lines changed: 101 additions & 43 deletions

File tree

databricks-skills/databricks-jobs/SKILL.md

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,43 @@ tasks:
218218
custom_param: "value"
219219
```
220220

221+
## MCP Tool Integration
222+
223+
Use these MCP tools for job management:
224+
225+
```python
226+
# Create or update a job
227+
manage_jobs(action="create", job_config={
228+
"name": "my_etl_job",
229+
"tasks": [{"task_key": "extract", "notebook_task": {"notebook_path": "/src/extract"}}]
230+
})
231+
232+
# List all jobs
233+
manage_jobs(action="list")
234+
235+
# Get job details
236+
manage_jobs(action="get", job_id=12345)
237+
238+
# Delete a job
239+
manage_jobs(action="delete", job_id=12345)
240+
241+
# Run a job immediately
242+
manage_job_runs(action="run_now", job_id=12345)
243+
244+
# Run with parameters
245+
manage_job_runs(action="run_now", job_id=12345,
246+
job_parameters={"env": "prod", "date": "2024-01-15"})
247+
248+
# Check run status
249+
manage_job_runs(action="get_run", run_id=67890)
250+
251+
# Cancel a run
252+
manage_job_runs(action="cancel", run_id=67890)
253+
254+
# List recent runs for a job
255+
manage_job_runs(action="list_runs", job_id=12345)
256+
```
257+
221258
## Common Operations
222259

223260
### Python SDK Operations

databricks-skills/databricks-unity-catalog/SKILL.md

Lines changed: 40 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,11 +85,48 @@ GROUP BY workspace_id, sku_name;
8585

8686
## MCP Tool Integration
8787

88-
Use `mcp__databricks__execute_sql` for system table queries:
88+
### Governance Tools
89+
90+
Use these MCP tools for Unity Catalog governance operations:
91+
92+
```python
93+
# Manage catalogs, schemas, tables, volumes, functions
94+
manage_uc_objects(action="list", object_type="catalogs")
95+
manage_uc_objects(action="create", object_type="schema", catalog="main", schema="my_schema")
96+
manage_uc_objects(action="describe", object_type="table", full_name="main.schema.table")
97+
98+
# Manage grants and permissions
99+
manage_uc_grants(action="list", securable_type="table", full_name="main.schema.table")
100+
manage_uc_grants(action="grant", securable_type="schema", full_name="main.my_schema",
101+
principal="data-engineers", privileges=["USE_SCHEMA", "SELECT"])
102+
103+
# Manage tags for classification
104+
manage_uc_tags(action="set", securable_type="table", full_name="main.schema.table",
105+
tags={"pii": "true", "team": "analytics"})
106+
107+
# Manage storage credentials and external locations
108+
manage_uc_storage(action="list", storage_type="credentials")
109+
manage_uc_storage(action="list", storage_type="external_locations")
110+
111+
# Manage Lakehouse Federation connections
112+
manage_uc_connections(action="list")
113+
114+
# Manage row filters and column masks
115+
manage_uc_security_policies(action="list", securable_type="table", full_name="main.schema.table")
116+
117+
# Manage data quality monitors
118+
manage_uc_monitors(action="list", table_name="main.schema.table")
119+
120+
# Manage Delta Sharing
121+
manage_uc_sharing(action="list", sharing_type="shares")
122+
```
123+
124+
### SQL Queries
125+
126+
Use `execute_sql` for system table queries:
89127

90128
```python
91-
# Query lineage
92-
mcp__databricks__execute_sql(
129+
execute_sql(
93130
sql_query="""
94131
SELECT source_table_full_name, target_table_full_name
95132
FROM system.access.table_lineage

databricks-tools-core/databricks_tools_core/auth.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,7 @@ def get_workspace_client() -> WorkspaceClient:
160160
# Cross-workspace: explicit token overrides env OAuth so tool operations
161161
# target the caller-specified workspace instead of the app's own workspace
162162
if force and host and token:
163-
return tag_client(
164-
WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs)
165-
)
163+
return tag_client(WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs))
166164

167165
# In Databricks Apps (OAuth credentials in env), explicitly use OAuth M2M.
168166
# Setting auth_type="oauth-m2m" prevents the SDK from also reading
@@ -185,9 +183,7 @@ def get_workspace_client() -> WorkspaceClient:
185183

186184
# Development mode: use explicit token if provided
187185
if host and token:
188-
return tag_client(
189-
WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs)
190-
)
186+
return tag_client(WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs))
191187

192188
if host:
193189
return tag_client(WorkspaceClient(host=host, **product_kwargs))

databricks-tools-core/tests/unit/test_sql.py

Lines changed: 21 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,7 @@ def test_executor_without_query_tags_omits_from_api(self, mock_get_client):
121121
assert "query_tags" not in call_kwargs
122122

123123

124-
def _make_warehouse(id, name, state, creator_name="other@example.com",
125-
enable_serverless_compute=False):
124+
def _make_warehouse(id, name, state, creator_name="other@example.com", enable_serverless_compute=False):
126125
"""Helper to create a mock warehouse object."""
127126
w = mock.Mock()
128127
w.id = id
@@ -141,33 +140,29 @@ class TestSortWithinTier:
141140
def test_serverless_first(self):
142141
"""Serverless warehouses should come before classic ones."""
143142
classic = _make_warehouse("c1", "Classic WH", State.RUNNING)
144-
serverless = _make_warehouse("s1", "Serverless WH", State.RUNNING,
145-
enable_serverless_compute=True)
143+
serverless = _make_warehouse("s1", "Serverless WH", State.RUNNING, enable_serverless_compute=True)
146144
result = _sort_within_tier([classic, serverless], current_user=None)
147145
assert result[0].id == "s1"
148146
assert result[1].id == "c1"
149147

150148
def test_serverless_before_user_owned(self):
151149
"""Serverless should be preferred over user-owned classic."""
152-
classic_owned = _make_warehouse("c1", "My WH", State.RUNNING,
153-
creator_name="me@example.com")
154-
serverless_other = _make_warehouse("s1", "Other WH", State.RUNNING,
155-
creator_name="other@example.com",
156-
enable_serverless_compute=True)
157-
result = _sort_within_tier([classic_owned, serverless_other],
158-
current_user="me@example.com")
150+
classic_owned = _make_warehouse("c1", "My WH", State.RUNNING, creator_name="me@example.com")
151+
serverless_other = _make_warehouse(
152+
"s1", "Other WH", State.RUNNING, creator_name="other@example.com", enable_serverless_compute=True
153+
)
154+
result = _sort_within_tier([classic_owned, serverless_other], current_user="me@example.com")
159155
assert result[0].id == "s1"
160156

161157
def test_serverless_user_owned_first(self):
162158
"""Among serverless, user-owned should come first."""
163-
serverless_other = _make_warehouse("s1", "Other Serverless", State.RUNNING,
164-
creator_name="other@example.com",
165-
enable_serverless_compute=True)
166-
serverless_owned = _make_warehouse("s2", "My Serverless", State.RUNNING,
167-
creator_name="me@example.com",
168-
enable_serverless_compute=True)
169-
result = _sort_within_tier([serverless_other, serverless_owned],
170-
current_user="me@example.com")
159+
serverless_other = _make_warehouse(
160+
"s1", "Other Serverless", State.RUNNING, creator_name="other@example.com", enable_serverless_compute=True
161+
)
162+
serverless_owned = _make_warehouse(
163+
"s2", "My Serverless", State.RUNNING, creator_name="me@example.com", enable_serverless_compute=True
164+
)
165+
result = _sort_within_tier([serverless_other, serverless_owned], current_user="me@example.com")
171166
assert result[0].id == "s2"
172167
assert result[1].id == "s1"
173168

@@ -177,53 +172,46 @@ def test_empty_list(self):
177172
def test_no_current_user(self):
178173
"""Without a current user, only serverless preference applies."""
179174
classic = _make_warehouse("c1", "Classic", State.RUNNING)
180-
serverless = _make_warehouse("s1", "Serverless", State.RUNNING,
181-
enable_serverless_compute=True)
175+
serverless = _make_warehouse("s1", "Serverless", State.RUNNING, enable_serverless_compute=True)
182176
result = _sort_within_tier([classic, serverless], current_user=None)
183177
assert result[0].id == "s1"
184178

185179

186180
class TestGetBestWarehouseServerless:
187181
"""Tests for serverless preference in get_best_warehouse."""
188182

189-
@mock.patch("databricks_tools_core.sql.warehouse.get_current_username",
190-
return_value="me@example.com")
183+
@mock.patch("databricks_tools_core.sql.warehouse.get_current_username", return_value="me@example.com")
191184
@mock.patch("databricks_tools_core.sql.warehouse.get_workspace_client")
192185
def test_prefers_serverless_within_running_shared(self, mock_client_fn, mock_user):
193186
"""Among running shared warehouses, serverless should be picked."""
194187
classic_shared = _make_warehouse("c1", "Shared WH", State.RUNNING)
195-
serverless_shared = _make_warehouse("s1", "Shared Serverless", State.RUNNING,
196-
enable_serverless_compute=True)
188+
serverless_shared = _make_warehouse("s1", "Shared Serverless", State.RUNNING, enable_serverless_compute=True)
197189
mock_client = mock.Mock()
198190
mock_client.warehouses.list.return_value = [classic_shared, serverless_shared]
199191
mock_client_fn.return_value = mock_client
200192

201193
result = get_best_warehouse()
202194
assert result == "s1"
203195

204-
@mock.patch("databricks_tools_core.sql.warehouse.get_current_username",
205-
return_value="me@example.com")
196+
@mock.patch("databricks_tools_core.sql.warehouse.get_current_username", return_value="me@example.com")
206197
@mock.patch("databricks_tools_core.sql.warehouse.get_workspace_client")
207198
def test_prefers_serverless_within_running_other(self, mock_client_fn, mock_user):
208199
"""Among running non-shared warehouses, serverless should be picked."""
209200
classic = _make_warehouse("c1", "My WH", State.RUNNING)
210-
serverless = _make_warehouse("s1", "Fast WH", State.RUNNING,
211-
enable_serverless_compute=True)
201+
serverless = _make_warehouse("s1", "Fast WH", State.RUNNING, enable_serverless_compute=True)
212202
mock_client = mock.Mock()
213203
mock_client.warehouses.list.return_value = [classic, serverless]
214204
mock_client_fn.return_value = mock_client
215205

216206
result = get_best_warehouse()
217207
assert result == "s1"
218208

219-
@mock.patch("databricks_tools_core.sql.warehouse.get_current_username",
220-
return_value="me@example.com")
209+
@mock.patch("databricks_tools_core.sql.warehouse.get_current_username", return_value="me@example.com")
221210
@mock.patch("databricks_tools_core.sql.warehouse.get_workspace_client")
222211
def test_tier_order_preserved_over_serverless(self, mock_client_fn, mock_user):
223212
"""A running shared classic should still beat a stopped serverless."""
224213
running_shared_classic = _make_warehouse("c1", "Shared WH", State.RUNNING)
225-
stopped_serverless = _make_warehouse("s1", "Fast WH", State.STOPPED,
226-
enable_serverless_compute=True)
214+
stopped_serverless = _make_warehouse("s1", "Fast WH", State.STOPPED, enable_serverless_compute=True)
227215
mock_client = mock.Mock()
228216
mock_client.warehouses.list.return_value = [stopped_serverless, running_shared_classic]
229217
mock_client_fn.return_value = mock_client

install.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -690,7 +690,7 @@ install_skills() {
690690
# Determine target directories (array so paths with spaces work)
691691
for tool in $TOOLS; do
692692
case $tool in
693-
claude) dirs=("$base_dir/.claude/skills") ;;
693+
claude) dirs+=("$base_dir/.claude/skills") ;;
694694
cursor) echo "$TOOLS" | grep -q claude || dirs+=("$base_dir/.cursor/skills") ;;
695695
copilot) dirs+=("$base_dir/.github/skills") ;;
696696
codex) dirs+=("$base_dir/.agents/skills") ;;

0 commit comments

Comments
 (0)