Skip to content

Commit 8f9475f

Browse files
smurchingsmurching
authored
Add initial UC MCP server integration (#2)
* WIP Signed-off-by: Sid Murching <[email protected]> * WIP adding tests Signed-off-by: Sid Murching <[email protected]> * Add fake tools and vibe check / stress test tool selection Signed-off-by: Sid Murching <[email protected]> * Update Signed-off-by: Sid Murching <[email protected]> * Add lint and unit tests Signed-off-by: Sid Murching <[email protected]> * fix Signed-off-by: Sid Murching <[email protected]> * Update runner for actions Signed-off-by: Sid Murching <[email protected]> * Add more labels Signed-off-by: Sid Murching <[email protected]> * fix Signed-off-by: Sid Murching <[email protected]> * Fix Signed-off-by: Sid Murching <[email protected]> * fix Signed-off-by: Sid Murching <[email protected]> * Fix Signed-off-by: Sid Murching <[email protected]> * Fix Signed-off-by: Sid Murching <[email protected]> * use official action for installing uv Signed-off-by: Sid Murching <[email protected]> * Fix lint Signed-off-by: Sid Murching <[email protected]> * Update lint script to have better error Signed-off-by: Sid Murching <[email protected]> * Update workflow name Signed-off-by: Sid Murching <[email protected]> * Fix trigger Signed-off-by: Sid Murching <[email protected]> * update Signed-off-by: Sid Murching <[email protected]> * Fix lint Signed-off-by: Sid Murching <[email protected]> * Clarify in readme Signed-off-by: Sid Murching <[email protected]> * Fix readme command Signed-off-by: Sid Murching <[email protected]> * Update readme Signed-off-by: Sid Murching <[email protected]> * Update READMEs and address comments Signed-off-by: Sid Murching <[email protected]> * Fix lint Signed-off-by: Sid Murching <[email protected]> * gs Signed-off-by: Sid Murching <[email protected]> * fix Signed-off-by: Sid Murching <[email protected]> --------- Signed-off-by: Sid Murching <[email protected]>
1 parent e1e7a91 commit 8f9475f

25 files changed

+4146
-19
lines changed

‎.github/workflows/unit_tests.yml

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
name: Tests and lint for MCP servers
2+
3+
on:
4+
push:
5+
paths:
6+
- "servers/**"
7+
- ".github/workflows/unit_tests.yml"
8+
pull_request:
9+
paths:
10+
- "servers/**"
11+
- ".github/workflows/unit_tests.yml"
12+
13+
jobs:
14+
discover:
15+
runs-on:
16+
group: databrickslabs-protected-runner-group
17+
labels: linux-ubuntu-latest
18+
outputs:
19+
server_dirs: ${{ steps.set-matrix.outputs.server_dirs }}
20+
21+
steps:
22+
- name: Checkout code
23+
uses: actions/checkout@v4
24+
25+
- name: Discover server dirs with tests
26+
id: set-matrix
27+
run: |
28+
dirs=$(find servers -mindepth 1 -maxdepth 1 -type d -exec test -d "{}/tests" \; -print | sed 's|servers/||' | jq -R . | jq -cs .)
29+
echo "Found directories: $dirs"
30+
echo "server_dirs=$dirs" >> "$GITHUB_OUTPUT"
31+
32+
33+
34+
test:
35+
needs: discover
36+
runs-on:
37+
group: databrickslabs-protected-runner-group
38+
labels: linux-ubuntu-latest
39+
strategy:
40+
matrix:
41+
server-dir: ${{ fromJson(needs.discover.outputs.server_dirs) }}
42+
43+
steps:
44+
- name: Checkout code
45+
uses: actions/checkout@v4
46+
47+
- name: Set up Python
48+
uses: actions/setup-python@v5
49+
with:
50+
python-version: "3.12"
51+
52+
- name: Install uv
53+
uses: astral-sh/setup-uv@v5
54+
55+
- name: Install dependencies with uv
56+
working-directory: servers/${{ matrix.server-dir }}
57+
run: |
58+
uv sync
59+
uv pip install --group dev
60+
61+
- name: Run tests
62+
working-directory: servers/${{ matrix.server-dir }}
63+
run: uv run pytest tests
64+
lint:
65+
runs-on:
66+
group: databrickslabs-protected-runner-group
67+
labels: linux-ubuntu-latest
68+
steps:
69+
- name: Checkout code
70+
uses: actions/checkout@v4
71+
72+
- name: Set up Python
73+
uses: actions/setup-python@v5
74+
with:
75+
python-version: "3.12"
76+
77+
- name: Install uv
78+
run: |
79+
curl -LsSf https://astral.sh/uv/install.sh | sh
80+
echo "$HOME/.cargo/bin" >> $GITHUB_PATH
81+
82+
- name: Run lint checks
83+
run: ./dev/lint.sh

‎.gitignore

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# MLflow
2+
mlruns/
3+
mlartifacts/
4+
outputs/
5+
mlruns.db
6+
**/basic_auth.db
7+
8+
# Mac
9+
.DS_Store
10+
11+
# Byte-compiled / optimized / DLL files
12+
__pycache__
13+
*.py[cod]
14+
*$py.class
15+
16+
# C extensions
17+
*.so
18+
19+
# Distribution / packaging
20+
.Python
21+
build/
22+
develop-eggs/
23+
dist/
24+
downloads/
25+
eggs/
26+
.eggs/
27+
lib/
28+
lib64/
29+
parts/
30+
sdist/
31+
var/
32+
wheels/
33+
*.egg-info/
34+
.installed.cfg
35+
*.egg
36+
MANIFEST
37+
node_modules
38+
.docusaurus
39+
*-ipynb.mdx
40+
docs/src/api_modules.json
41+
docs/static/api_reference
42+
43+
# PyInstaller
44+
# Usually these files are written by a python script from a template
45+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
46+
*.manifest
47+
*.spec
48+
49+
# Installer logs
50+
pip-log.txt
51+
pip-delete-this-directory.txt
52+
53+
# Unit test / coverage reports
54+
htmlcov/
55+
.coverage
56+
.coverage.*
57+
.cache
58+
nosetests.xml
59+
coverage.xml
60+
*.cover
61+
.hypothesis/
62+
.pytest_cache/
63+
64+
# Sphinx documentation
65+
docs/_build/
66+
67+
# Jupyter Notebook
68+
.ipynb_checkpoints
69+
70+
# Environments
71+
env
72+
env3
73+
.env
74+
.venv
75+
env/
76+
venv/
77+
ENV/
78+
env.bak/
79+
venv.bak/
80+
.python-version
81+
dev-env-setup-progress
82+
83+
# Editor files
84+
.*project
85+
*.swp
86+
*.swo
87+
*.idea
88+
*.vscode
89+
*.iml
90+
*~
91+
92+
# mkdocs documentation
93+
/site
94+
95+
# mypy
96+
.mypy_cache/
97+
98+
# java targets
99+
target/
100+
101+
# R notebooks
102+
.Rproj.user
103+
example/tutorial/R/*.nb.html
104+
105+
# travis_wait command logs
106+
travis_wait*.log
107+
108+
# Pytorch logs
109+
lightning_logs
110+
111+
a.py
112+
a.ipynb
113+
a.md
114+
115+
# Log file created by pre-commit hook for black
116+
.black.log
117+
118+
# Pytest-monitor load testing DB file
119+
*.pymon
120+
121+
# Ignore a gunicorn config file
122+
gunicorn.conf.py

‎CONTRIBUTING.md

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,26 @@ We happily welcome contributions to the MCP servers in this repo. We use GitHub
55
For new feature requests or contributions (e.g. adding a new server), please file an issue or reach out [via Slack](https://databricks.enterprise.slack.com/archives/C065NC65Q9F) to facilitate initial discussion,
66
before sending a pull request
77

8+
### Running tests
9+
10+
First, install test requirements from within the directory of the server you're working on:
11+
12+
```bash
13+
uv sync
14+
uv pip install --group dev
15+
```
16+
17+
To run tests:
18+
19+
```bash
20+
uv run pytest tests
21+
```
22+
23+
To run and fix lint errors, run the following from the repo root directory:
24+
```bash
25+
./dev/lint.sh --fix
26+
```
27+
828
### Guidelines for MCP servers
929

1030
For consistency, MCP servers in this repo:

‎README.md

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
# Databricks MCP servers
2-
![status: WIP](https://img.shields.io/badge/status-WIP-red?style=flat-square&logo=databricks)
3-
4-
## 🚧 Work in Progress 🚧
5-
6-
**This repo is still under initial development and must not be shared outside Databricks.**
2+
![status: Beta](https://img.shields.io/badge/status-Beta-yellow?style=flat-square&logo=databricks)
73

84
## Overview
9-
A collection of [MCP](https://modelcontextprotocol.io/introduction) servers to help AI agents take common developer actions on Databricks, fetch data from Databricks, etc:
5+
An experimental collection of [MCP](https://modelcontextprotocol.io/introduction) servers to help AI agents fetch enterprise data from Databricks, automate common developer actions on Databricks, etc:
6+
7+
* [Databricks Unity Catalog server](./servers/unity_catalog/README.md): Fetch data and run tools registered in from Unity Catalog, making agents aware of your enterprise data
8+
* 🚧 (Work in progress) [Databricks developer tools server](./servers/developer_tools/README.md): Perform common developer actions in Databricks, like creating and updating notebooks, running jobs, etc.
109

11-
* 🚧 [Databricks Unity Catalog server](./unity_catalog): Fetch data and run tools registered in from Unity Catalog, making agents aware of your enterprise data
12-
* 🚧 [Databricks developer tools server](./developer_tools): Perform common developer actions in Databricks, like creating and updating notebooks, running jobs, etc.
10+
The set of servers in this repo is fluid and will evolve over time. We welcome contributions to this repo, including new servers and new tools for existing servers - please first
11+
read the [contributor guidelines](CONTRIBUTING.md) to streamline the process and ensure your contribution has the
12+
best change of being accepted :)
1313

1414
## Usage
1515
See the `README.md` in each server's directory for detailed instructions.
@@ -24,20 +24,24 @@ For most servers, the following steps work:
2424
{
2525
"mcpServers": {
2626
"databricks_unity_catalog": {
27-
"command": "uv",
27+
"command": "/path/to/uv/executable/uv",
2828
"args": [
2929
"--directory",
3030
"/path/to/this/repo/servers/unity_catalog",
31+
"run",
32+
"unitycatalog-mcp",
3133
"-s",
32-
"prod.genai"
34+
"your_catalog.your_schema",
35+
"-g",
36+
"genie_space_id_1,genie_space_id_2"
3337
]
3438
}
3539
}
3640
}
3741
```
3842

3943
## Support
40-
Please note that all projects in the /databrickslabs github account are provided for your exploration only, and are not formally supported by Databricks with Service Level Agreements (SLAs). They are provided AS-IS and we do not make any guarantees of any kind. Please do not submit a support ticket relating to any issues arising from the use of these projects.
44+
Please note that all projects in the `databrickslabs` GitHub organization are provided for your exploration only, and are not formally supported by Databricks with Service Level Agreements (SLAs). They are provided AS-IS and we do not make any guarantees of any kind. Please do not submit a support ticket relating to any issues arising from the use of these projects.
4145

4246
Any issues discovered through the use of this project should be filed as GitHub Issues on the Repo. They will be reviewed as time permits, but there are no formal SLAs for support.
4347

‎dev/lint.sh

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#!/usr/bin/env bash
2+
set -euo pipefail
3+
4+
# Usage: ./lint.sh [--fix]
5+
6+
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
7+
FIX=false
8+
9+
if [[ "${1:-}" == "--fix" ]]; then
10+
FIX=true
11+
fi
12+
13+
echo "Linting Python files in servers/... with black and ruff"
14+
uv pip install --system black ruff
15+
16+
if $FIX; then
17+
echo "🛠 Fixing formatting with black and ruff"
18+
black servers/
19+
ruff check servers/ --fix
20+
else
21+
echo "🔍 Checking formatting (without fixing)"
22+
if ! black servers/ --check; then
23+
echo -e "\n❌ Black found formatting issues. To fix them, run:"
24+
echo " ./dev/lint.sh --fix"
25+
exit 1
26+
fi
27+
28+
if ! ruff check servers/; then
29+
echo -e "\n❌ Ruff found issues. To fix them, run:"
30+
echo " ./dev/lint.sh --fix"
31+
exit 1
32+
fi
33+
fi

‎servers/developer_tools/README.md

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
![status: WIP](https://img.shields.io/badge/status-WIP-red?style=flat-square&logo=databricks)
33

44
## 🚧 Work in Progress 🚧
5-
**This server is still under initial development and must not be shared outside Databricks.**
5+
**This server is still under initial development.**
66

77
## Overview
88
A Model Context Protocol server that exposes common Databricks developer actions as tools.
@@ -16,11 +16,13 @@ A Model Context Protocol server that exposes common Databricks developer actions
1616
```json
1717
{
1818
"mcpServers": {
19-
"databricks_unity_catalog": {
20-
"command": "uv",
19+
"databricks_developer_tools": {
20+
"command": "/path/to/uv/executable/uv",
2121
"args": [
2222
"--directory",
23-
"/path/to/this/repo/servers/developer_tools"
23+
"/path/to/this/repo/servers/developer_tools",
24+
"run",
25+
"databricks-developer-tools"
2426
]
2527
}
2628
}

‎servers/unity_catalog/LICENSE

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# DB license
2+
Definitions.
3+
4+
Agreement: The agreement between Databricks, Inc., and you governing the use of the Databricks Services, as that term is defined in the Master Cloud Services Agreement (MCSA) located at www.databricks.com/legal/mcsa.
5+
6+
Licensed Materials: The source code, object code, data, and/or other works to which this license applies.
7+
8+
Scope of Use. You may not use the Licensed Materials except in connection with your use of the Databricks Services pursuant to the Agreement. Your use of the Licensed Materials must comply at all times with any restrictions applicable to the Databricks Services, generally, and must be used in accordance with any applicable documentation. You may view, use, copy, modify, publish, and/or distribute the Licensed Materials solely for the purposes of using the Licensed Materials within or connecting to the Databricks Services. If you do not agree to these terms, you may not view, use, copy, modify, publish, and/or distribute the Licensed Materials.
9+
10+
Redistribution. You may redistribute and sublicense the Licensed Materials so long as all use is in compliance with these terms. In addition:
11+
12+
You must give any other recipients a copy of this License;
13+
You must cause any modified files to carry prominent notices stating that you changed the files;
14+
You must retain, in any derivative works that you distribute, all copyright, patent, trademark, and attribution notices, excluding those notices that do not pertain to any part of the derivative works; and
15+
If a "NOTICE" text file is provided as part of its distribution, then any derivative works that you distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the derivative works.
16+
You may add your own copyright statement to your modifications and may provide additional license terms and conditions for use, reproduction, or distribution of your modifications, or for any such derivative works as a whole, provided your use, reproduction, and distribution of the Licensed Materials otherwise complies with the conditions stated in this License.
17+
18+
Termination. This license terminates automatically upon your breach of these terms or upon the termination of your Agreement. Additionally, Databricks may terminate this license at any time on notice. Upon termination, you must permanently delete the Licensed Materials and all copies thereof.
19+
20+
21+
DISCLAIMER; LIMITATION OF LIABILITY.
22+
23+
THE LICENSED MATERIALS ARE PROVIDED “AS-IS” AND WITH ALL FAULTS. DATABRICKS, ON BEHALF OF ITSELF AND ITS LICENSORS, SPECIFICALLY DISCLAIMS ALL WARRANTIES RELATING TO THE LICENSED MATERIALS, EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, IMPLIED WARRANTIES, CONDITIONS AND OTHER TERMS OF MERCHANTABILITY, SATISFACTORY QUALITY OR FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT. DATABRICKS AND ITS LICENSORS TOTAL AGGREGATE LIABILITY RELATING TO OR ARISING OUT OF YOUR USE OF OR DATABRICKS’ PROVISIONING OF THE LICENSED MATERIALS SHALL BE LIMITED TO ONE THOUSAND ($1,000) DOLLARS. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE LICENSED MATERIALS OR THE USE OR OTHER DEALINGS IN THE LICENSED MATERIALS.

‎servers/unity_catalog/README.md

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,16 @@ Genie spaces, as tools.
1818
{
1919
"mcpServers": {
2020
"databricks_unity_catalog": {
21-
"command": "uv",
21+
"command": "/path/to/uv/executable/uv",
2222
"args": [
2323
"--directory",
2424
"/path/to/this/repo/servers/unity_catalog",
25+
"run",
26+
"unitycatalog-mcp",
2527
"-s",
26-
// Replace with the name of your Unity Catalog schema
27-
"prod.genai"
28+
"your_catalog.your_schema",
29+
"-g",
30+
"genie_space_id_1,genie_space_id_2"
2831
]
2932
}
3033
}
@@ -33,7 +36,7 @@ Genie spaces, as tools.
3336

3437
## Supported tools
3538

36-
The list of tools supported by this server is dynamically inferred based on the functions and vector search indexes
39+
The list of tools supported by this server is dynamically inferred at startup time based on the functions and vector search indexes
3740
within the specified Unity Catalog schema, as well as any specified Genie spaces. In particular, the server exposes
3841
the following tools:
3942

0 commit comments

Comments
 (0)