Databricks MCP Server

Databricks MCP 服务

一个模型上下文协议（MCP）服务，可连接到 Databricks API，让大模型能够运行 SQL 查询、列出作业并获取作业状态。

功能

• 在 Databricks SQL 仓库中运行 SQL 查询
• 列出所有 Databricks 作业
• 获取特定 Databricks 作业的状态
• 获取 Databricks 作业的详细信息

前提条件

• Python 3.7+
• Databricks 工作区，包含：
  • 个人访问令牌
  • SQL 仓库端点
  • 运行查询和访问作业的权限

设置

1. 克隆此代码仓库
2. 创建并激活虚拟环境（推荐）：

   python -m venv .venv
   source .venv/bin/activate  # On Windows: .venv\Scripts\activate
   

3. 安装依赖：

   pip install -r requirements.txt
   

4. 在根目录中创建 .env 文件，并设置以下变量：

   DATABRICKS_HOST=your-databricks-instance.cloud.databricks.com
   DATABRICKS_TOKEN=your-personal-access-token
   DATABRICKS_HTTP_PATH=/sql/1.0/warehouses/your-warehouse-id
   

5. 测试连接（可选，但推荐）：

   python test_connection.py
   

获取 Databricks 凭据

1. 主机：你的 Databricks 实例 URL（例如，your-instance.cloud.databricks.com）
2. 令牌：在 Databricks 中创建个人访问令牌：
   • 前往用户设置（点击右上角的用户名）
   • 选择“Developer”选项卡
   • 点击“Access tokens”下的“Manage”
   • 生成新令牌并立即保存
3. HTTP 路径：针对你的 SQL 仓库：
   • 在 Databricks 中前往 SQL 仓库
   • 选择你的仓库
   • 找到连接详情并复制 HTTP 路径

运行服务

启动 MCP 服务：

python main.py

你可以运行以下命令，通过检查器测试 MCP 服务：

npx @modelcontextprotocol/inspector python3 main.py

可用的 MCP 工具

以下 MCP 工具可用：

1. run_sql_query(sql: str) - 在你的 Databricks SQL 仓库中执行 SQL 查询
2. list_jobs() - 列出工作区中的所有 Databricks 作业
3. get_job_status(job_id: int) - 根据 ID 获取特定 Databricks 作业的状态
4. get_job_details(job_id: int) - 获取特定 Databricks 作业的详细信息

使用大模型的示例

与支持 MCP 协议的大模型配合使用时，此服务可让你通过自然语言与 Databricks 环境交互：

• “显示数据库中的所有表”
• “运行查询，统计 customer 表中的记录数”
• “列出我的所有 Databricks 作业”
• “检查作业 #123 的状态”
• “显示作业 #456 的详细信息”

故障排除

连接问题

• 确保 Databricks 主机地址正确，且不包含 https:// 前缀
• 检查 SQL 仓库是否正在运行且可以访问
• 确认个人访问令牌具有所需权限
• 运行附带的测试脚本：python test_connection.py

安全注意事项

• Databricks 个人访问令牌可直接访问你的工作区
• 妥善保护 .env 文件，切勿将其提交到版本控制系统
• 考虑仅使用具有适当权限范围的 Databricks 令牌
• 在安全的环境中运行此服务