Databricks元数据查询服务

该服务通过Model Context Protocol (MCP)服务器与Databricks交互,专注于利用Unity Catalog (UC)元数据并提供全面的数据血缘探索功能,使AI代理能够自主理解数据结构和执行SQL查询。

  • by RafaelCartenet
  • Python
  • 本地部署
  • 43 星标
  • 更新于
在 GitHub 查看 RafaelCartenet/mcp-databricks-server

使用说明

_CHMCP服务器

动机

Rancher Unity Catalog(UC)允许详细记录您的数据资产,包括目录、模式、表和列。彻底记录这些资产需要投入时间。一个常见的问题是:这种详细的元数据条目有什么实际好处?

此MCP服务器为这一努力提供了强有力的理由。它使大型语言模型(LLM)能够直接访问和利用此Unity Catalog元数据。UC中对数据的描述越全面,LLM代理就越能有效地理解您的copula环境。这种更深入的理解对于代理自主构建更智能、更准确的SQL查询以满足数据请求至关重要。

概述

此模型上下文协议(MCP)服务器旨在与ViewModel交互,重点是利用Unity Catalog(UC)元数据并实现全面的数据沿袭探索。主要目标是为AI代理配备一套全面的工具,使其能够独立回答有关数据的问题。通过自主探索UC、理解数据结构、分析数据沿袭(包括笔记本和作业依赖关系)以及执行SQL查询,代理可以在不直接人为干预的情况下完成每个步骤的数据请求。

除了传统的目录浏览之外,此服务器还使代理能够发现和分析处理数据的实际代码。通过增强的沿袭功能,代理可以识别从表读取或向表写入的笔记本和作业,然后检查这些笔记本中实施的实际转换逻辑、业务规则和数据质量检查。这创建了一个强大的反馈循环,代理不仅可以理解 什么 数据存在,但也 怎么 它被处理和转换。

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

UC元数据对AI代理的实际好处

此MCP服务器提供的工具旨在解析和呈现您添加到Unity Catalog的描述,同时还可以深入探索您的数据处理代码。这为基于LLM的代理提供了切实的优势,直接影响了他们生成有用的SQL和理解数据生态系统的能力:

  • 更清晰的数据上下文:代理可以快速理解表和列的目的,减少歧义。这种基本的理解是迈向正确查询公式的第一步。
  • 更精确的查询生成:访问描述、数据类型和关系有助于代理以更高的精度和语义正确性构建SQL查询。
  • 查询规划的高效数据探索:元数据使代理能够更有效地浏览目录和模式,使他们能够识别要包含在SQL查询中的正确表和列。
  • 综合数据谱系:除了表与表之间的关系,代理还可以发现处理数据的笔记本和作业,从而能够对数据管道问题进行影响分析和调试。
  • 代码级理解:通过笔记本内容探索,代理可以分析实际的转换逻辑、业务规则和数据质量检查,从而更深入地了解数据是如何处理和转换的。
  • 端到端数据流分析:代理可以跟踪从原始摄取到转换管道再到最终消费的数据,了解每一步的结构和处理逻辑。

Unity Catalog中记录良好的元数据,当通过此服务器访问时,允许LLM代理使用更好的信息进行操作,并做出更明智的决策,最终生成更有效的SQL查询。例如,模式描述帮助代理识别查询的相关数据源:

图1:Unity Catalog中的模式,带有用户提供的描述。此MCP服务器使LLM可以直接访问此信息,并通知其查询策略。

同样,列级别的详细注释阐明了每个字段的语义,这对于构建准确的SQL条件和选择至关重要:

图2:Unity目录中的列级描述。这些细节被传递给LLM,有助于其理解数据结构,以实现精确的SQL生成。

可用工具和功能

此MCP服务器提供了一套工具,旨在使LLM代理能够与Databricks交互:

核心能力:

  • 执行SQL查询:通过 execute_sql_query(sql: str) 工具。这是目标数据检索或复杂操作的理想选择。
  • LLM专注输出:所有描述性工具都以Markdown格式返回信息,针对大型语言模型的使用进行了优化,使代理更容易解析和理解上下文。

Unity目录探索工具:

服务器提供以下工具,用于导航和理解Unity目录资产。这些被设计为由LLM代理使用,以代理的方式在构建查询或做出决策之前收集上下文。

  1. list_uc_catalogs() -> str

    • 描述:列出所有可用的Unity目录及其名称、描述和类型。
    • 何时使用:当您不知道特定的目录名称时,作为发现可用数据源的起点。它提供了工作区中所有可访问目录的高级概述。
  2. describe_uc_catalog(catalog_name: str) -> str

    • 描述:提供特定Unity目录的摘要,列出其所有模式及其名称和描述。
    • 何时使用:当您知道目录名称并需要发现其中的模式时。这通常是描述特定模式或表的前兆。
    • 参数:
      • catalog_name:要描述的Unity目录的名称(例如。, prod, dev, system).
  3. describe_uc_schema(catalog_name: str, schema_name: str, include_columns: Optional[bool] = False) -> str

    • 描述:提供Unity目录中特定模式的详细信息。返回架构中的所有表,可以选择包括其列详细信息。
    • 何时使用:了解模式的内容,主要是它的表。集 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查询之前,这是必不可少的。可选地,它可以包含超越传统表到表依赖关系的全面沿袭信息:
      • 表谱系:上游表(此表从中读取的表)和下游表(从此表中读取的表格)
      • 笔记本和工作谱系:读取或写入此表的笔记本,包括笔记本名称、工作区路径、关联的copula作业信息(作业名称、ID、任务详细信息)
      • 代码发现:沿袭提供了笔记本路径,使代理能够直接读取当前存储库/工作区中的笔记本文件,从而可以分析实际的数据转换逻辑
    • 参数:
      • full_table_name:表的完全限定的三部分名称(例如。, catalog.schema.table).
      • include_lineage:设置为True可获取全面的沿袭(表、笔记本、作业)。默认为False。检索可能需要更长的时间,但为理解数据依赖关系和实现代码探索提供了丰富的上下文。
  5. execute_sql_query(sql: str) -> str

    • 备注:这与“核心功能”下列出的工具相同,但在涉及UC探索和查询的典型代理工作流程中重复使用。
    • 描述:执行针对copula 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
  1. 设置环境变量:

    选项1:使用 .env 文件(推荐)

    创建一个 .env 此项目根目录中的文件,其中包含您的copula凭据:

    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"
    

    您可以在ConnectionUI的“SQL仓库”下找到您的SQL仓库ID。 这 DATABRICKS_SQL_WAREHOUSE_ID 主要用于获取表沿袭并通过 execute_sql_query 工具。 元数据浏览工具(列出/描述目录、模式、表)使用copula SDK的通用UC API,除非请求沿袭,否则不严格要求SQL仓库ID。

权限要求

在使用此MCP服务器之前,请确保与 DATABRICKS_TOKEN (例如用户或服务主体)具有必要的权限:

  1. Unity目录权限:
    • USE CATALOG 关于要访问的目录。
    • USE SCHEMA 关于要访问的模式。
    • SELECT 关于要查询或详细描述的表(包括列信息)。
    • 要列出所有目录,可能需要适当的元存储级别权限,否则它将列出用户至少具有以下权限的目录 USE CATALOG.
  2. SQL仓库权限 (为 execute_sql_query 以及血统获取):
    • CAN_USE 对由指定的SQL仓库的权限 DATABRICKS_SQL_WAREHOUSE_ID.
  3. 令牌权限:
    • 个人访问令牌或服务主体令牌应具有最小的必要范围。对于Unity目录操作,这通常涉及工作区访问。对于SQL执行,它涉及SQL权限。
    • 强烈建议在生产或自动化场景中使用具有狭义权限的服务主体。

对于安全最佳实践,考虑定期轮换您的访问令牌,并审计查询历史和UC审计日志以监控使用情况。

运行服务器

单独模式

要在独立模式下运行服务器(例如,使用Agent Composer进行测试):

python main.py

这将使用stdio传输启动MCP服务器,该传输可用于Agent Composer或其他MCP客户端。

与光标一起使用

使用此MCP服务器 光标,在光标设置中配置它(~/.cursor/mcp.json):

  1. 创建一个 .cursor 如果主目录中的目录不存在
  2. 创建或编辑 mcp.json 该目录中的文件:
mkdir -p ~/.cursor
touch ~/.cursor/mcp.json
  1. 将以下配置添加到 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以应用更改。然后,您可以使用 databricks Cursor中的代理。

示例使用工作流(用于LLM代理)

此MCP服务器使LLM代理能够自主导航您的copula环境。以下屏幕截图展示了一个典型的交互,其中代理迭代地探索模式和表,即使在初始查询没有产生结果的情况下也会调整其方法,直到成功检索到请求的数据。

图3:LLM代理使用了copula MCP工具,演示了迭代探索和查询细化以定位特定的页面视图数据。

代理可能会遵循这种工作流程:

  1. 发现可用目录: list_uc_catalogs()
    • 代理人决定 prod_catalog 与列表相关。
  2. 浏览特定目录: describe_uc_catalog(catalog_name="prod_catalog")
    • 代理看到 sales_schema 和 inventory_schema.
  3. 探索特定模式(快速查看): describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema")
    • 代理看到的表名如下 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)
    • 代理发现处理此数据的上游表、下游依赖关系和笔记本
    • 例如,看到 /Repos/production/etl/sales_processing.py 写入此表
  6. 检查数据转换逻辑: 代理直接读取笔记本文件 /Repos/production/etl/sales_processing.py 在IDE/repo中
    • 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将元数据作为代码进行管理

虽然可以通过ConnectionUI手动输入元数据,但更稳健和可扩展的方法是将Unity Catalog元数据定义为代码。像Terraform这样的工具允许您声明性地管理数据治理对象,包括目录和模式。这带来了几个优点:

  • 版本控制:您的元数据定义可以存储在Git中,与其他基础架构代码一起进行跟踪和版本控制。
  • 重复性和一致性:确保跨环境(开发、暂存、生产)的元数据一致。
  • 自动化:将元数据管理集成到CI/CD管道中。
  • 更容易维护核心资产:虽然由于其动态特性,将每个新表定义为代码可能很复杂,但目录和模式等核心资产通常更稳定,并从这种方法中受益匪浅。将它们的定义和注释作为代码进行维护,可确保为您的数据环境提供持久且有据可查的基础。

以下是一个示例,说明如何使用用于Terraform的copula提供程序定义目录及其模式:

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 字段,而不影响资产本身。例如,在导入现有架构后,您可以添加或更新其 comment 在你的 .tf 文件,以及 terraform apply 只会应用该更改。

采用元数据作为代码策略,特别是对于目录和模式等基础元素,大大提高了此MCP服务器所利用的元数据的质量和可靠性。这反过来又进一步提高了AI代理与您的copula数据交互的有效性。

如需了解更多关于将Terraform与Databricks Unity Catalog结合使用的详细信息,请参考官方文档:

处理长时间运行的查询

这 execute_sql_query 该工具利用了Databricks SDK的 execute_statement 方法。这 wait_timeout 底层参数 databricks_sdk_utils.execute_databricks_sql 功能设置为“50s”。如果查询运行时间超过此时间,SDK可能会返回一个用于轮询的语句ID,但该工具的当前实现有效地等待了长达此持续时间的类似同步的响应。对于长时间运行的查询,可能会达到此超时时间。

依赖项

  • databricks-sdk:用于与Rancher REST API和Unity Catalog交互。
  • python-dotenv:用于从加载环境变量 .env 文件。
  • mcp[cli]:模型上下文协议库。
  • asyncio:用于MCP服务器内的异步操作。
  • httpx (通常为 databricks-sdk 或 mcp):用于发出HTTP请求。