Databricks Genie API MCP Server

Databricks Genie API MCP 服务

本项目实现了一个 Model Context Protocol（MCP）服务，将 Databricks Genie API 功能作为工具提供。它可通过标准化接口，将 Databricks 的无代码 AI/BI 助手功能与其他应用集成，从而使用自然语言强大地查询 Databricks 数据。

如需详细说明和示例用例，请参阅随附的博客文章。

功能

• 将 Databricks Genie API 函数作为 MCP 工具提供。
• 支持使用自然语言查询 Databricks 数据。
• 启动和管理 Genie 对话。
• 创建和检索消息。
• 执行并获取 Genie 生成的 SQL 查询结果。
• 使用 Databricks 进行安全身份验证。

前提条件

• Python 3.10+
• 已启用 Genie 访问权限的 Databricks 工作区；如需使用 System Tables，也需启用该功能。
• 已启用 Databricks Assistant。
• 对 Pro 或 Serverless SQL 仓库具有 CAN USE 权限。
• 可访问与 Genie Space 相关的 Unity Catalog 数据。
• 兼容 MCP 的客户端应用，例如 Claude Desktop。

设置说明

1. 克隆代码库（如果尚未克隆）

   # Add clone command if needed
   

2. 进入服务目录

   cd genie_api 
   

3. 安装依赖项

   pip install -r requirements.txt
   

4. 配置身份验证

   设置 Databricks SDK 连接到工作区所需的环境变量。在此目录（genie_api/）中创建 .env 文件，或设置为全局环境变量：

   # --- .env file content ---
   
   # For PAT authentication (recommended for development)
   # DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
   # DATABRICKS_TOKEN=your-personal-access-token
   
   # Or for OAuth with service principal (recommended for production)
   DATABRICKS_HOST=https://your-workspace.cloud.databricks.com
   DATABRICKS_CLIENT_ID=your-client-id
   DATABRICKS_CLIENT_SECRET=your-client-secret
   
   # --- end .env file content ---
   

   请确保将 .env 文件添加到 .gitignore 中！

5. 在本地运行服务

   python server.py
   

   服务将启动，并通过标准输入/输出（stdio）监听连接。

与 Claude Desktop 搭配使用

此 MCP 服务旨在与 Claude Desktop 等 MCP 客户端搭配使用。请按照以下步骤进行连接：

1. 安装 Claude Desktop： 从官方网站下载并安装。

2. 配置 Claude Desktop：

   • 打开 Claude Desktop 设置（菜单栏 -> Claude -> Settings...）。
   • 前往 Developer -> Edit Config。
   • 在 claude_desktop_config.json 文件的 mcpServers 对象中添加以下条目，并根据需要调整路径：

   {
     "mcpServers": {
       "databricks-genie": {
         "command": "python", // Or python3, or the full path to your python executable
         "args": [
           "/full/absolute/path/to/your/project/genie_api/server.py" 
         ],
         "workingDirectory": "/full/absolute/path/to/your/project/genie_api/" 
       }
       // ... potentially other servers ...
     }
   }
   

   • 重要： 使用 server.py 和 genie_api 目录的完整绝对路径。
   • 确保 Claude Desktop 可以访问 Python command。
   • workingDirectory 可确保服务能够找到 auth.py 和你的 .env 文件。

3. 重启 Claude Desktop： 关闭并重新打开应用。

4. 验证： 点击聊天输入框中的锤子图标（工具）。你应该能看到列出的 databricks-genie 工具（例如 start_conversation、create_message）。

现在，你可以向 Claude 提问，例如“我们上个月消耗了多少 DBU？”或“昨天谁访问了 PII 表？”，Claude 会使用本地服务提供的工具。

有关配置 Claude Desktop 的更多详情，请参阅面向 Claude Desktop 用户的 MCP 快速入门指南。

可用工具

• start_conversation：在 Genie 空间中开始新的对话。
• create_message：在现有对话中创建新消息。
• get_message：从对话中获取消息。
• get_message_attachment_query_result：获取消息附件中的 SQL 查询结果。
• execute_message_attachment_query：执行消息查询附件中的 SQL。
• get_space：获取 Genie 空间的详细信息。
• generate_download_full_query_result：启动完整查询结果的下载。
• poll_message_until_complete：轮询消息，直到其达到终止状态。

故障排除

• 身份验证问题：检查 Databricks 凭据（.env 文件或环境变量）以及 Databricks 中所需的权限。
• 连接问题（Claude Desktop）：
  • 确认 claude_desktop_config.json 中的绝对路径正确。
  • 检查 command 是否指向有效的 Python 解释器。
  • 检查 Claude Desktop 日志（~/Library/Logs/Claude/ 或 %APPDATA%\Claude\logs）。查找 mcp.log 和 mcp-server-databricks-genie.log。
  • 尝试在终端中从 genie_api 目录手动运行 python server.py，以检查是否有错误。
• SQL 执行错误：确保服务主体或用户对 SQL 仓库拥有 CAN USE 权限，并且可以访问相关的 Unity Catalog 数据。

安全注意事项

• **凭据：**切勿将凭据硬编码。请使用环境变量（.env）或安全的凭据存储。确保将 .env 添加到 .gitignore 中。生产环境中请使用服务主体进行 OAuth 认证。
• MCP 安全性：此服务在本地运行，使用用户的权限和 Databricks 凭据。Claude Desktop 等客户端在执行工具前必须征得用户同意（MCP 安全规范）。
• **生产环境部署：**若要供更多人使用，必须制定安全的托管方案。**切勿直接公开此本地服务。**请与安全/DevOps 团队合作，确定合适的托管方式、网络控制措施，以及是否需要 MCP 服务级身份验证。托管环境必须能够安全地访问 Databricks 凭据（例如，实例配置文件、托管密钥）。更多讨论请参阅博客文章。
• **输入清理：**对于通过工具传入的输入，请信任 Databricks SDK/API 进行处理。

贡献

欢迎贡献！请提交 issue 或拉取请求。

许可证

本软件依据 Databricks, Inc. 授予的特定许可证提供。有关规范您使用本软件的完整条款与条件，请参阅 LICENSE 文件。