Databricks MCP Server

Databricks MCP 服务

• 动机
• 概述
• Unity Catalog 元数据对 AI Agent 的实际益处
• 可用工具和功能
• 设置
  • 系统要求
  • 安装
• 权限要求
• 运行服务
  • 独立模式
  • 与 Cursor 配合使用
• 示例使用工作流（针对大模型 Agent）
• 使用 Terraform 将元数据作为代码管理
• 处理长时间运行的查询
• 依赖项

动机

Databricks Unity Catalog（UC）允许为数据资产编写详细文档，包括目录、架构、表和列。要全面记录这些资产，需要投入时间。一个常见问题是：详细录入这些元数据有什么实际好处？

此 MCP 服务为这项工作提供了充分的理由。它使大语言模型（LLM）能够直接访问并利用 Unity Catalog 元数据。UC 中对数据的描述越全面，LLM Agent 就越能有效地理解你的 Databricks 环境。更深入的理解对于 Agent 自主构建更智能、更准确的 SQL 查询以满足数据需求至关重要。

概述

此模型上下文协议（MCP）服务旨在与 Databricks 交互，重点利用 Unity Catalog（UC）元数据并支持全面的数据血缘探索。其主要目标是为 AI Agent 配备一套全面的工具，使其能够独立回答有关数据的问题。通过自主探索 UC、理解数据结构、分析数据血缘（包括笔记本和作业依赖关系）以及执行 SQL 查询，Agent 无需每一步都由人工直接介入，即可满足数据需求。

除了传统的目录浏览功能，此服务还允许 Agent 发现并分析处理数据的实际代码。借助增强的数据血缘功能，Agent 可以识别读取或写入表的笔记本和作业，然后检查这些笔记本中实现的实际转换逻辑、业务规则和数据质量检查。这形成了一个强大的反馈闭环，让 Agent 不仅理解数据 是什么，还理解数据 如何 处理和转换。

在 Agent 模式下使用时，它可以成功迭代处理多个请求，执行数据发现、影响分析和代码探索等复杂任务。

UC 元数据对 AI Agent 的实际优势

此 MCP 服务提供的工具旨在解析并呈现你添加到 Unity Catalog 的描述，同时支持深入探索数据处理代码。这为基于大模型的 Agent 带来了切实优势，直接提升了它们生成实用 SQL 和理解数据生态系统的能力：

• 更清晰的数据上下文：Agent 可以快速了解表和列的用途，减少歧义。这种基础认知是正确构造查询的第一步。
• 更准确地生成查询：访问描述、数据类型和关系，有助于 Agent 更精准地构造语义正确的 SQL 查询。
• 高效探索数据以规划查询：元数据使 Agent 能更有效地浏览目录和架构，从而确定 SQL 查询中应包含的表和列。
• 全面的数据沿袭：除了表与表之间的关系，Agent 还可以发现处理数据的笔记本和作业，从而分析影响并调试数据管道问题。
• 代码级理解：通过探索笔记本内容，Agent 可以分析实际的转换逻辑、业务规则和数据质量检查，从而深入了解数据的处理和转换方式。
• 端到端数据流分析：Agent 可以追踪数据从原始摄取到转换管道，再到最终使用的全过程，了解每个步骤中的结构和处理逻辑。

通过此服务访问 Unity Catalog 中记录完善的元数据，能让大模型 Agent 掌握更充分的信息并做出更明智的决策，最终生成更有效的 SQL 查询。例如，架构描述可以帮助 Agent 找到与查询相关的数据源：

图 1：Unity Catalog 中包含用户提供描述的架构。此 MCP 服务可让大模型直接访问这些信息，为其查询策略提供依据。

同样，列级别的详细注释可以明确各字段的语义，这对于构造准确的 SQL 条件和选择项至关重要：

图 2：Unity Catalog 中的列级别描述。这些细节会传递给大模型，帮助其理解数据结构，从而精准生成 SQL。

可用工具和功能

此 MCP 服务提供了一套工具，旨在帮助与 Databricks 交互的大模型 Agent：

核心功能：

• 执行 SQL 查询：通过 execute_sql_query(sql: str) 工具使用 Databricks SDK 运行任意 SQL 查询。适用于定向检索数据或执行复杂操作。
• 面向大模型的输出：所有描述类工具均以 Markdown 格式返回信息，并针对大模型的使用进行了优化，让 Agent 更轻松地解析并理解上下文。

Unity Catalog 探索工具：

服务提供以下工具，用于浏览和了解你的 Unity Catalog 资产。这些工具旨在供大模型 Agent 以 Agentic 方式使用，在构造查询或做出决策之前收集上下文。

1. list_uc_catalogs() -> str

   • 说明：列出所有可用的 Unity Catalog，并提供其名称、描述和类型。
   • 使用时机：当你不知道具体的目录名称时，可从此工具开始发现可用的数据源。它会概览工作区中所有可访问的目录。

2. describe_uc_catalog(catalog_name: str) -> str

   • 描述：提供特定 Unity Catalog 的摘要，列出其中所有架构的名称和描述。
   • 使用场景：已知目录名称，并需要查找其中的架构时。这通常是描述特定架构或表的前置步骤。
   • 参数：
     • catalog_name：要描述的 Unity Catalog 名称（例如，prod、dev、system）。

3. describe_uc_schema(catalog_name: str, schema_name: str, include_columns: Optional[bool] = False) -> str

   • 描述：提供 Unity Catalog 中特定架构的详细信息。返回该架构中的所有表，并可选择包含其列的详细信息。
   • 使用场景：了解架构的内容，主要是其中的表。设置 include_columns=True 可获取列信息，这对构造查询至关重要，但会使输出更长。如果 include_columns=False，则只显示表名和描述，便于快速浏览。
   • 参数：
     • catalog_name：包含该架构的目录名称。
     • schema_name：要描述的架构名称。
     • include_columns：如果为 True，则列出表及其列。默认为 False，以提供更简洁的摘要。

4. describe_uc_table(full_table_name: str, include_lineage: Optional[bool] = False) -> str

   • 描述：提供特定 Unity Catalog 表的详细说明，并具备全面的血缘分析能力。
   • 使用场景：了解单个表的结构（列、数据类型、分区）。在针对该表构造 SQL 查询之前，这一步至关重要。此外，还可以选择包含全面的血缘信息，涵盖传统表间依赖关系以外的内容：
     • 表血缘：上游表（该表读取数据的来源）和下游表（读取该表数据的表）
     • Notebook 和作业血缘：读取或写入该表的 Notebook，包括 Notebook 名称、工作区路径、关联的 Databricks 作业信息（作业名称、ID、任务详细信息）
     • 代码发现：血缘信息会提供 Notebook 路径，使 Agent 能够直接读取当前仓库/工作区中的 Notebook 文件，从而分析实际的数据转换逻辑
   • 参数：
     • full_table_name：表的完整三段式名称（例如，catalog.schema.table）。
     • include_lineage：设置为 True 可获取全面的血缘信息（表、Notebook、作业）。默认为 False。检索可能需要更长时间，但能提供丰富的上下文，帮助理解数据依赖关系并探索代码。

5. execute_sql_query(sql: str) -> str

   • 注意：这是“核心功能”中列出的同一工具，此处再次列出，以说明典型 Agent 工作流中先探索 UC，再执行查询的场景。
   • 描述：对 Databricks SQL 仓库执行指定的 SQL 查询，并返回格式化结果。
   • 使用场景：需要运行特定 SQL 查询时，例如 SELECT、SHOW 或其他 DQL 语句。
   • 参数：
     • sql：要执行的完整 SQL 查询字符串。

设置

系统要求

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

安装

1. 安装所需的依赖项：

pip install -r requirements.txt

或者使用 uv：

uv pip install -r requirements.txt

2. 设置环境变量：

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

   在此项目的根目录中创建一个 .env 文件，并填入 Databricks 凭据：

   DATABRICKS_HOST="your-databricks-instance.cloud.databricks.com"
   DATABRICKS_TOKEN="your-databricks-personal-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-personal-access-token"
   export DATABRICKS_SQL_WAREHOUSE_ID="your-sql-warehouse-id"
   

   你可以在 Databricks UI 的“SQL Warehouses”下找到 SQL Warehouse ID。 DATABRICKS_SQL_WAREHOUSE_ID 主要用于获取表的血缘信息，以及通过 execute_sql_query 工具执行 SQL 查询。 元数据浏览工具（列出／描述目录、架构和表）使用 Databricks SDK 的通用 UC API，除非需要查询血缘信息，否则不严格要求提供 SQL Warehouse ID。

权限要求

使用此 MCP 服务前，请确保与 DATABRICKS_TOKEN 关联的身份（例如用户或服务主体）具有所需权限：

1. Unity Catalog 权限：
   • 对要访问的目录具有 USE CATALOG 权限。
   • 对要访问的架构具有 USE SCHEMA 权限。
   • 对要查询或详细描述（包括列信息）的表具有 SELECT 权限。
   • 要列出所有目录，可能需要适当的 metastore 级权限；否则，只会列出该用户至少具有 USE CATALOG 权限的目录。
2. SQL Warehouse 权限（用于 execute_sql_query 和获取血缘信息）：
   • 对 DATABRICKS_SQL_WAREHOUSE_ID 指定的 SQL Warehouse 具有 CAN_USE 权限。
3. 令牌权限：
   • 个人访问令牌或服务主体令牌应具有必要的最小权限范围。对于 Unity Catalog 操作，通常需要工作区访问权限；对于 SQL 执行，则需要 SQL 权限。
   • 强烈建议在生产或自动化场景中使用权限范围严格受限的服务主体。

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

运行服务

独立模式

要以独立模式运行服务（例如，用于 Agent Composer 测试）：

python main.py

这将通过 stdio 传输启动 MCP 服务，可与 Agent Composer 或其他 MCP 客户端配合使用。

与 Cursor 配合使用

要将此 MCP 服务与 Cursor 配合使用，请在 Cursor 设置（~/.cursor/mcp.json）中进行配置：

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"
            ]
        }
    }
}

使用 python 的示例：

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

重启 Cursor 以应用更改。之后，你就可以在 Cursor 中使用 databricks Agent。

示例使用工作流（适用于大模型 Agent）

此 MCP 服务可让大模型 Agent 自主探索你的 Databricks 环境。以下截图展示了典型的交互过程：Agent 逐步探索架构和表，即使初始查询没有结果，也会调整探索方式，直到成功获取所需数据。

图 3：大模型 Agent 使用 Databricks MCP 工具，通过反复探索和优化查询来查找特定的页面浏览数据。

Agent 可能会遵循以下工作流：

1. 发现可用目录：list_uc_catalogs()
   • Agent 根据列表判断 prod_catalog 相关。
2. 探索特定目录：describe_uc_catalog(catalog_name="prod_catalog")
   • Agent 看到 sales_schema 和 inventory_schema。
3. 探索特定架构（快速查看）：describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema")
   • Agent 看到 orders、customers 等表名。
4. 获取详细的表结构（包括用于构建查询的列）：describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema", include_columns=True)
   • 或者，如果关注的是某个特定表： describe_uc_table(full_table_name="prod_catalog.sales_schema.orders")
5. 分析数据沿袭并查找处理代码：describe_uc_table(full_table_name="prod_catalog.sales_schema.orders", include_lineage=True)
   • Agent 会发现上游表、下游依赖项，以及处理这些数据的笔记本。
   • 例如，Agent 会看到 /Repos/production/etl/sales_processing.py 会向此表写入数据。
6. 检查数据转换逻辑：Agent 直接在 IDE/代码仓库中读取笔记本文件 /Repos/production/etl/sales_processing.py。
   • Agent 分析实际的 Python/SQL 代码，以了解业务规则、数据质量检查和转换逻辑。
7. 构建并执行查询：execute_sql_query(sql="SELECT customer_id, order_date, SUM(order_total) FROM prod_catalog.sales_schema.orders WHERE order_date > '2023-01-01' GROUP BY customer_id, order_date ORDER BY order_date DESC LIMIT 100")

使用 Terraform 将元数据作为代码进行管理

虽然可以通过 Databricks UI 手动输入元数据，但更稳健且可扩展的做法是将 Unity Catalog 元数据定义为代码。Terraform 等工具支持以声明式方式管理数据治理对象，包括目录和架构。这样做有以下优势：

• 版本控制：可以将元数据定义存储在 Git 中，与其他基础设施代码一起跟踪和管理版本。
• 可重复性和一致性：确保各个环境（开发、预发布、生产）中的元数据保持一致。
• 自动化：将元数据管理集成到 CI/CD 管道中。
• 更轻松地维护核心资产：由于新表具有动态特性，将每个新表都定义为代码可能较为复杂，但目录和架构等核心资产通常更加稳定，因此能从这种做法中显著受益。将这些资产的定义和注释作为代码维护，可为数据环境奠定持久且文档完善的基础。

以下示例展示了如何使用 Terraform 的 Databricks 提供程序定义目录及其架构：

resource "databricks_catalog" "prod_catalog" {
  name          = "prod"
  comment       = "Main production catalog for all enterprise data."
  storage_root  = var.default_catalog_storage_root
  force_destroy = false
}

# Schemas within the 'prod' catalog
resource "databricks_schema" "prod_raw" {
  catalog_name = databricks_catalog.prod_catalog.name
  name         = "raw"
  comment      = "Raw data for all different projects, telemetry, game data etc., before any transformations. No schema enforcement."
}

resource "databricks_schema" "prod_bi_conformed" {
  catalog_name = databricks_catalog.prod_catalog.name
  name         = "bi_conformed"
  comment      = "Conformed (silver) schema for Business Intelligence, cleaned and well-formatted. Schema enforced."
}

resource "databricks_schema" "prod_bi_modeled" {
  catalog_name = databricks_catalog.prod_catalog.name
  name         = "bi_modeled"
  comment      = "Modeled (gold) schema for Business Intelligence, aggregated and ready for consumption. Schema enforced."
}

如果 Unity Catalog 中已有目录和架构，也不必担心。无需重新创建它们即可将其元数据作为代码进行管理。Terraform 提供了 terraform import 命令，可将现有基础设施（包括 Unity Catalog 资产）纳入其管理。导入后，便可在 Terraform 配置中定义该资源，并有选择地更新 comment 字段等属性，而不会影响资产本身。例如，导入现有架构后，可以在 .tf 文件中添加或更新其 comment，然后运行 terraform apply，只应用这项更改。

采用元数据即代码策略，尤其是针对目录和架构等基础元素，可大幅提升此 MCP 服务所用元数据的质量和可靠性，进而进一步提升 AI Agent 与 Databricks 数据交互的效果。

有关在 Databricks Unity Catalog 中使用 Terraform 的更多详情，请参阅官方文档：

• Databricks 提供程序：目录资源（https://registry.terraform.io/providers/databricks/databricks/latest/docs/resources/catalog）
• Databricks 提供程序：架构数据源（https://registry.terraform.io/providers/databricks/databricks/latest/docs/data-sources/schemas）

处理长时间运行的查询

execute_sql_query 工具使用 Databricks SDK 的 execute_statement 方法。底层 databricks_sdk_utils.execute_databricks_sql 函数中的 wait_timeout 参数设为 '50s'。如果查询运行时间超过此时长，SDK 可能会返回一个用于轮询的语句 ID，但该工具的当前实现实际上最多会等待此时长，以获得类似同步调用的响应。对于运行时间很长的查询，可能会达到此超时限制。

依赖项

• databricks-sdk：用于与 Databricks REST API 和 Unity Catalog 交互。
• python-dotenv：用于从 .env 文件加载环境变量。
• mcp[cli]：Model Context Protocol 库。
• asyncio：用于 MCP 服务中的异步操作。
• httpx（通常是 databricks-sdk 或 mcp 的间接依赖项）：用于发送 HTTP 请求。