MCP Server for Databricks Interaction

Databricks 交互 MCP 服务

1. 项目目标

本项目提供一个服务器应用，用作 Databricks 工作区的接口。它使用 FastMCP 框架公开工具，让用户能够直接通过 MCP 命令查询和检索 Databricks 资源的信息，例如架构、表、样本和作业结果。

主要目标是借助 Databricks SDK 和 CLI 进行后端通信，简化用户通过 MCP 接口执行常见 Databricks 元数据检索任务的流程。

服务的主要目标和功能：

此模型上下文协议旨在提供可与 IDE 配合使用的工具，为你的 Agent 提供 Databricks 环境上下文（尽管这些工具并非只对基于 IDE 的开发工作有用）。这样一来，在协助你为工作区编写代码时，Agent 的效率会得到提升，你需要手动提供的额外信息也会减少。其架构如下：

1. OAuth U2M 身份验证：此服务使用 Databricks cli 在你与 Databricks 工作区之间进行身份验证。这样就无需创建令牌并将其存储在本地，而大多数替代方案都需要这样做。
2. 提供目录对象的元数据：现有 MCP 工具使用 Databricks SDK 提供目录、架构和表的元数据。同样，你也可以获取工作流的运行结果，以便进行改进或自行排查问题。
3. 提供表的样本数据：使用 get_table_sample_tool 检索表相关信息时，此服务还会使用用户可访问的 SQL 仓库检索样本行，然后将这些样本行添加到每列的元数据中，以丰富所提供的上下文。

2. 安装和初始化

前提条件

• Python（推荐使用 3.x 版本）
• uv（Python 包安装器和虚拟环境管理器）。你可以按照 https://github.com/astral-sh/uv 上的说明安装。
• 已安装 databricks-cli，且可通过系统的 PATH 访问。

重要说明： pypi 上提供的 databricks_cli 版本为 0.18，版本相当旧，本项目不支持。要安装最新版本的 databricks_cli，请访问 Databricks 官方页面并按照说明操作。

安装

1. 克隆代码库：

   git clone <repository_url>
   cd <repository_directory>
   

2. 创建虚拟环境并安装依赖项：

   uv venv  # 创建虚拟环境（例如 .venv）
   uv sync  # 使用 pyproject.toml 和 uv.lock 安装依赖项
   source .venv/bin/activate # 激活虚拟环境（在 Windows 上使用 `.venv\\Scripts\\activate`）
   

初始化

首次运行服务前，你需要配置它与 Databricks 工作区的连接：

1. 运行初始化脚本：

   python init.py
   

2. 按照提示操作：

   • 输入 Databricks 工作区 URL（例如，https://adb-xxxxxxxxxxxx.azuredatabricks.net）。
   • 脚本将使用配置文件名 mcp_server_for_databricks 启动 Databricks CLI OAuth 登录流程。按照 CLI 在终端/浏览器中提供的说明完成身份验证。
   • 从提供的列表中选择一个可用的 SQL Warehouse。可用的 MCP 工具将使用此 Warehouse 执行元数据查询。
   • 输入使用表采样工具时要检索的样本大小（行数）。按 Enter 使用默认值（5）。

3. 保存配置： 脚本会将工作区 URL、所选 Warehouse ID/名称和样本大小保存到项目根目录下的 config.yaml 文件中。

4. （适用于 Cursor）添加到 Cursor IDE

初始化服务并安装依赖项后，你可以将其添加到 Cursor：

1. 打开 Cursor 设置： 前往 Cursor Settings > Features > MCP。
2. 添加新的 MCP 服务： 点击 + Add New MCP Server 按钮。
3. 配置服务： 按照下面的示例粘贴 MCP 服务详细信息

{
  "mcpServers": {
    "mcp_server_for_databricks": {
      "command": "/path/to/uv/executable/uv",
      "args": [
        "--directory",
        "/PATH/TO/THIS/REPO",
        "run",
        "main.py"
      ]
    }
  }
}

4. 保存并刷新： 保存配置。你可能需要点击 MCP 列表中该服务旁的刷新按钮，才能加载其工具。

3. 工作原理

• 服务框架： 使用 FastMCP 定义工具，并通过标准输入/输出（stdio）处理通信。
• Databricks 连接： 使用 Python 的 databricks-sdk 与 Databricks 交互。
• 身份验证：
  • 依赖名为 mcp_server_for_databricks 的 databricks-cli 配置文件。
  • main.py 脚本使用 subprocess 调用 databricks auth login 和 databricks auth token，以获取并为 SDK 客户端使用身份验证令牌。身份验证会在服务首次启动时进行；如果服务尚未初始化，则会在首次调用工具时进行。
• 信息检索：
  • 元数据查询： 大多数架构和表列表操作都会通过 client.statement_execution.execute_statement SDK 方法，在已配置的 SQL Warehouse 上执行 SQL 语句（SHOW DATABASES、SHOW TABLES）。
  • 详细表信息： 使用 client.tables.get SDK 方法获取完整的表详细信息（列类型、注释、属性等）。
  • 表采样： 通过 client.statement_execution.execute_statement 方法执行 SELECT * ... LIMIT N 查询，以检索样本数据。
  • 作业结果： 使用 client.jobs API（list、list_runs、get_run、get_run_output）查找作业并检索特定运行的输出。
• 配置： 启动时从 config.yaml 读取连接详细信息（工作区 URL、Warehouse ID、样本大小）。

4. 提供的 MCP 工具

此服务提供以下工具：

• get_schemas(catalog: str)

  • 描述： 检索指定目录中所有架构及其包含的表的列表。
  • 参数：
    • catalog（str）：要查询的 Databricks 目录名称。
  • 返回： 对象列表，每个对象包含 catalog、schema_name 和一个 tables 列表。

• get_table_sample_tool(catalog: str, schema_name: str, table: str)

  • **说明：**返回特定表的详细元数据，包括列信息（类型、注释、是否可为空）和表数据中的示例行。示例行数由初始化时配置的 sample_size 决定。如果在 config.yaml 中配置，还会将元数据和示例数据保存到 .input_data/ 目录。
  • 参数：
    • catalog (str)：目录名称。
    • schema_name (str)：架构名称。
    • table (str)：表名称。
  • **返回值：**返回的字典结构如下：

    {
        "name": "string (表名称)",
        "catalog_name": "string (目录名称)",
        "schema_name": "string (架构名称)",
        "table_type": "string (例如，MANAGED、EXTERNAL、VIEW)",
        "data_source_format": "string (例如，DELTA、CSV、PARQUET)",
        "columns": [
            {
                "name": "string (列名称)",
                "type_name": "string (列数据类型，例如 INT、STRING)",
                "comment": "string (列注释或 null)",
                "nullable": "boolean",
                "partition_index": "integer (或 null；如果该列是分区列，则为从 0 开始的索引)",
                "mask": "object (列掩码详情或 null)",
                "sample_values": ["示例值列表（混合类型）"]
            }
            // ... 更多列
        ],
        "comment": "string (表注释或 null)",
        "properties": {
            "property_key": "property_value"
            // ... 更多属性
        },
        "storage_location": "string (存储路径；对于托管表/视图为 null)",
        "view_definition": "string (如果是视图则为 SQL 定义，否则为 null)",
        "table_id": "string (唯一表标识符)",
        "created_at": "string (ISO 8601 时间戳或 null)",
        "updated_at": "string (ISO 8601 时间戳或 null)",
        "deleted_at": "string (ISO 8601 时间戳或 null，如适用)",
        "row_filter": "object (行筛选器详情或 null)",
        "owner": "string (表所有者，例如用户或组)"
    }
    

• get_schema_metadata(catalog_name: str, schema_name: str)

  • **说明：**获取特定架构的元数据，包括架构注释以及其中每个表的详情（表注释、创建时间戳、表类型、所有者）。
  • 参数：
    • catalog_name (str)：目录名称。
    • schema_name (str)：架构名称。
  • **返回值：**返回的字典结构如下（tables 字典将包含该架构中每个表对应的条目）：

    {
        "schema_comment": "string (架构注释或 null)",
        "tables": {
            "your_table_name_here": {
                "comment": "string (表注释或 null)",
                "created_at": "string (ISO 8601 时间戳)",
                "table_type": "string (例如，MANAGED、EXTERNAL、VIEW)",
                "owner": "string (表所有者)"
            }
            // ... 更多表
        }
    }
    

• get_job_run_result(job_name: str, filter_for_failed_runs: bool = False)

  • 描述： 获取指定 Databricks 作业最近一次运行的结果。可选择获取最近一次运行的结果（无论其状态如何），或专门获取最近一次失败的运行结果。
  • 参数：
    • job_name (str)：Databricks 作业的准确名称。
    • filter_for_failed_runs (bool, optional)：如果为 True，则仅获取最近一次失败运行的结果。默认为 False（获取最近一次已完成运行的结果）。
  • 返回值： 一个字符串，包含所选作业运行相关的错误消息、错误堆栈跟踪和元数据。

已知问题：

• 目前，JWT 令牌超过有效期限后不会自动刷新。因此，如果 MCP 服务已运行超过 12 小时，就需要重启。此问题很快会得到修复
• 对于列数很多或注释较多的表，get_table_sample_tool 可能会检索过多上下文。如果用户遇到此问题，可以减小配置文件中的 sample_size，以降低此风险。
• 当前版本要求用户在需要从其他 Databricks 工作区检索信息时手动切换工作区，这对尚未迁移到 Unity Catalog 的组织来说是个重大问题。后续版本将支持指定不同的工作区