Databricks MCP 服务器

这是一个使用 Statement Execution API 对 Databricks 执行 SQL 查询的 Model Context Protocol (MCP) 服务器。它可以通过使用 Databricks API 执行 SQL 请求来检索数据。在代理模式下使用时，它可以成功迭代多个请求以执行复杂任务。与 Unity Catalog 元数据结合使用时效果更佳。

功能

• 在 Databricks 上执行 SQL 查询
• 列出目录中的可用模式
• 列出模式中的表
• 描述表结构

设置

系统要求

• Python 3.10+
• 如果您计划通过 uv 安装，请确保已安装

安装

1. 安装所需的依赖项：

pip install -r requirements.txt

或者如果使用 uv：

uv pip install -r requirements.txt

2. 设置环境变量：

   选项 1：使用 .env 文件（推荐）

   创建一个包含您的 Databricks 凭证的 .env 文件：

   DATABRICKS_HOST=your-databricks-instance.cloud.databricks.com
   DATABRICKS_TOKEN=your-databricks-access-token
   DATABRICKS_SQL_WAREHOUSE_ID=your-sql-warehouse-id
   

   选项 2：直接设置环境变量

   export DATABRICKS_HOST="your-databricks-instance.cloud.databricks.com"
   export DATABRICKS_TOKEN="your-databricks-access-token"
   export DATABRICKS_SQL_WAREHOUSE_ID="your-sql-warehouse-id"
   

您可以在 Databricks UI 中的 SQL Warehouses 下找到您的 SQL 仓库 ID。

权限要求

在使用此 MCP 服务器之前，请确保：

1. SQL 仓库权限：与所提供令牌关联的用户必须具有访问指定 SQL 仓库的适当权限。您可以在 Databricks UI 的 SQL Warehouses > [Your Warehouse] > Permissions 下配置仓库权限。

2. 令牌权限：使用的个人访问令牌应具有执行所需操作的最小必要权限。强烈建议：

   • 为此应用程序创建专用令牌
   • 尽可能授予只读权限以限制安全风险
   • 避免使用具有工作区范围管理员权限的令牌

3. 数据访问权限：与令牌关联的用户必须具有访问将要查询的目录、模式和表的适当权限。

要通过 Databricks REST API 设置 SQL 仓库权限，您可以使用：

• GET /api/2.0/sql/permissions/warehouses/{warehouse_id} 检查当前权限
• PATCH /api/2.0/sql/permissions/warehouses/{warehouse_id} 更新权限

为了遵循安全最佳实践，考虑定期轮换访问令牌并审核查询历史记录以监控使用情况。

运行服务器

独立模式

要在独立模式下运行服务器：

python main.py

这将使用 stdio 传输启动 MCP 服务器，可以与 Agent Composer 或其他 MCP 客户端一起使用。

与 Cursor 一起使用

要将此 MCP 服务器与 Cursor 一起使用，你需要在 Cursor 设置中进行配置：

1. 如果尚不存在，请在你的主目录下创建一个 .cursor 目录
2. 在该目录中创建或编辑 mcp.json 文件：

mkdir -p ~/.cursor
touch ~/.cursor/mcp.json

3. 将以下配置添加到 mcp.json 文件中，用你安装此服务器的实际路径替换目录路径：

{
    "mcpServers": {
        "databricks": {
            "command": "uv",
            "args": [
                "--directory",
                "/path/to/your/mcp-databricks-server",
                "run",
                "main.py"
            ]
        }
    }
}

如果你不使用 uv，可以改用 python：

{
    "mcpServers": {
        "databricks": {
            "command": "python",
            "args": [
                "/path/to/your/mcp-databricks-server/main.py"
            ]
        }
    }
}

4. 重启 Cursor 以应用更改

现在你可以在 Cursor 的 AI 助手中直接使用 Databricks MCP 服务器了。

可用工具

服务器提供以下工具：

1. execute_sql_query: 执行 SQL 查询并返回结果

   execute_sql_query(sql: str) -> str
   

2. list_schemas: 列出特定目录中的所有可用模式

   list_schemas(catalog: str) -> str
   

3. list_tables: 列出特定模式中的所有表

   list_tables(schema: str) -> str
   

4. describe_table: 描述表的结构

   describe_table(table_name: str) -> str
   

使用示例

在 Agent Composer 或其他 MCP 客户端中，你可以这样使用这些工具：

execute_sql_query("SELECT * FROM my_schema.my_table LIMIT 10")
list_schemas("my_catalog")
list_tables("my_catalog.my_schema")
describe_table("my_catalog.my_schema.my_table")

处理长时间运行的查询

服务器设计为通过轮询 Databricks API 来处理长时间运行的查询，直到查询完成或超时。默认超时时间为 10 分钟（60 次重试，每次间隔 10 秒），如果需要，可以在 dbapi.py 文件中调整。

依赖项

• httpx: 用于向 Databricks API 发送 HTTP 请求
• python-dotenv: 用于从 .env 文件加载环境变量
• mcp: Model Context Protocol 库
• asyncio: 用于异步操作