Databricks MCP Server

Interaja com os metadados do Databricks Unity Catalog para explorar catálogos, esquemas, tabelas e executar consultas SQL.

GitHub
42
Experimente este MCPPatrocinado

Documentação

Listed on Spark Install via Spark

Databricks MCP Server

Motivação

O Databricks Unity Catalog (UC) permite a documentação detalhada dos seus ativos de dados, incluindo catálogos, esquemas, tabelas e colunas. Documentar esses ativos minuciosamente exige um investimento de tempo. Uma pergunta comum é: quais são os benefícios práticos dessa entrada detalhada de metadados?

Este servidor MCP fornece uma forte justificativa para esse esforço. Ele permite que Modelos de Linguagem de Grande Porte (LLMs) acessem e utilizem diretamente esses metadados do Unity Catalog. Quanto mais abrangentemente seus dados forem descritos no UC, mais eficazmente um agente LLM poderá entender seu ambiente Databricks. Esse entendimento mais profundo é crucial para que o agente construa autonomamente consultas SQL mais inteligentes e precisas para atender às solicitações de dados.

Visão Geral

Este servidor Model Context Protocol (MCP) é projetado para interagir com o Databricks, com um forte foco em aproveitar os metadados do Unity Catalog (UC) e permitir uma exploração abrangente da linhagem de dados. O objetivo principal é equipar um agente de IA com um conjunto abrangente de ferramentas, permitindo que ele se torne independente ao responder perguntas sobre seus dados. Ao explorar autonomamente o UC, entender estruturas de dados, analisar a linhagem de dados (incluindo dependências de notebooks e jobs) e executar consultas SQL, o agente pode atender às solicitações de dados sem intervenção humana direta em cada etapa.

Além da navegação tradicional de catálogo, este servidor permite que os agentes descubram e analisem o código real que processa seus dados. Por meio de recursos aprimorados de linhagem, os agentes podem identificar notebooks e jobs que leem ou gravam em tabelas e, em seguida, examinar a lógica real de transformação, regras de negócios e verificações de qualidade de dados implementadas nesses notebooks. Isso cria um ciclo de feedback poderoso em que os agentes não apenas entendem o que dados existem, mas também como eles são processados e transformados.

Quando usado em um modo de Agente, ele pode iterar com sucesso sobre uma série de solicitações para executar tarefas complexas, incluindo descoberta de dados, análise de impacto e exploração de código.

Benefícios Práticos dos Metadados do UC para Agentes de IA

As ferramentas fornecidas por este servidor MCP são projetadas para analisar e apresentar as descrições que você adicionou ao Unity Catalog, além de permitir uma exploração profunda do seu código de processamento de dados. Isso oferece vantagens tangíveis para agentes baseados em LLM, impactando diretamente sua capacidade de gerar SQL útil e entender seu ecossistema de dados:

  • Contexto de Dados Mais Claro: Os agentes podem entender rapidamente a finalidade de tabelas e colunas, reduzindo a ambiguidade. Esse entendimento fundamental é o primeiro passo para a formulação correta de consultas.
  • Geração de Consultas Mais Precisa: O acesso a descrições, tipos de dados e relacionamentos ajuda os agentes a construir consultas SQL com maior precisão e correção semântica.
  • Exploração Eficiente de Dados para Planejamento de Consultas: Os metadados permitem que os agentes naveguem por catálogos e esquemas de forma mais eficaz, permitindo que identifiquem as tabelas e colunas corretas para incluir em suas consultas SQL.
  • Linhagem de Dados Abrangente: Além dos relacionamentos tabela a tabela, os agentes podem descobrir notebooks e jobs que processam dados, permitindo análise de impacto e depuração de problemas em pipelines de dados.
  • Entendimento em Nível de Código: Por meio da exploração do conteúdo dos notebooks, os agentes podem analisar a lógica real de transformação, regras de negócios e verificações de qualidade de dados, fornecendo insights mais profundos sobre como os dados são processados e transformados.
  • Análise de Fluxo de Dados de Ponta a Ponta: Os agentes podem rastrear dados desde a ingestão bruta, passando por pipelines de transformação, até o consumo final, entendendo tanto a estrutura quanto a lógica de processamento em cada etapa.

Metadados bem documentados no Unity Catalog, quando acessados por meio deste servidor, permitem que um agente LLM opere com melhores informações e tome decisões mais informadas, culminando na geração de consultas SQL mais eficazes. Por exemplo, descrições de esquemas ajudam o agente a identificar fontes de dados relevantes para uma consulta:

Schema Description in Unity Catalog Fig 1: Um esquema no Unity Catalog com descrições fornecidas pelo usuário. Este servidor MCP torna essas informações diretamente acessíveis a um LLM, informando sua estratégia de consulta.

Da mesma forma, comentários detalhados no nível da coluna esclarecem a semântica de cada campo, o que é crucial para construir condições e seleções SQL precisas:

Table Column Descriptions in Unity Catalog Fig 2: Descrições em nível de coluna no Unity Catalog. Esses detalhes são passados ao LLM, auxiliando seu entendimento da estrutura de dados para geração precisa de SQL.

Ferramentas e Recursos Disponíveis

Este servidor MCP fornece um conjunto de ferramentas projetadas para capacitar um agente LLM que interage com o Databricks:

Capacidades Principais:

  • Executar Consultas SQL: Execute consultas SQL arbitrárias usando o SDK do Databricks por meio da ferramenta execute_sql_query(sql: str). Isso é ideal para recuperação direcionada de dados ou operações complexas.
  • Saída Focada em LLM: Todas as ferramentas descritivas retornam informações em formato Markdown, otimizado para consumo por Modelos de Linguagem de Grande Porte, facilitando a análise e o entendimento do contexto pelos agentes.

Ferramentas de Exploração do Unity Catalog:

O servidor fornece as seguintes ferramentas para navegar e entender seus ativos do Unity Catalog. Elas são projetadas para serem usadas por um agente LLM para reunir contexto antes de construir consultas ou tomar decisões, de forma agêntica.

  1. list_uc_catalogs() -> str

    • Descrição: Lista todos os Unity Catalogs disponíveis com seus nomes, descrições e tipos.
    • Quando usar: Como ponto de partida para descobrir fontes de dados disponíveis quando você não conhece nomes específicos de catálogos. Fornece uma visão geral de alto nível de todos os catálogos acessíveis no workspace.

  2. describe_uc_catalog(catalog_name: str) -> str

    • Descrição: Fornece um resumo de um Unity Catalog específico, listando todos os seus esquemas com nomes e descrições.
    • Quando usar: Quando você conhece o nome do catálogo e precisa descobrir os esquemas dentro dele. Isso geralmente é um precursor para descrever um esquema ou tabela específica.
    • Argumentos:
      • catalog_name: O nome do Unity Catalog a ser descrito (por exemplo, prod, dev, system).

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

    • Descrição: Fornece informações detalhadas sobre um esquema específico dentro de um Unity Catalog. Retorna todas as tabelas no esquema, opcionalmente incluindo detalhes de suas colunas.
    • Quando usar: Para entender o conteúdo de um esquema, principalmente suas tabelas. Defina include_columns=True para obter informações de colunas, cruciais para a construção de consultas, mas torna a saída mais longa. Se include_columns=False, apenas nomes e descrições de tabelas são mostrados, útil para uma visão geral mais rápida.
    • Argumentos:
      • catalog_name: O nome do catálogo que contém o esquema.
      • schema_name: O nome do esquema a ser descrito.
      • include_columns: Se True, lista tabelas com suas colunas. Padrão é False para um resumo mais breve.

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

    • Descrição: Fornece uma descrição detalhada de uma tabela específica do Unity Catalog com capacidades abrangentes de linhagem.
    • Quando usar: Para entender a estrutura (colunas, tipos de dados, particionamento) de uma única tabela. Isso é essencial antes de construir consultas SQL contra a tabela. Opcionalmente, pode incluir informações abrangentes de linhagem que vão além das dependências tradicionais tabela a tabela:
      • Linhagem de Tabela: Tabelas upstream (tabelas das quais esta tabela lê) e tabelas downstream (tabelas que leem desta tabela)
      • Linhagem de Notebooks e Jobs: Notebooks que leem ou gravam nesta tabela, incluindo nome do notebook, caminho no workspace, informações associadas do job do Databricks (nome do job, ID, detalhes da tarefa)
      • Descoberta de Código: A linhagem fornece caminhos de notebooks que permitem que um agente leia diretamente arquivos de notebook no repositório/workspace atual, permitindo a análise da lógica real de transformação de dados
    • Argumentos:
      • full_table_name: O nome totalmente qualificado em três partes da tabela (por exemplo, catalog.schema.table).
      • include_lineage: Defina como True para buscar linhagem abrangente (tabelas, notebooks, jobs). Padrão é False. Pode levar mais tempo para recuperar, mas fornece contexto rico para entender dependências de dados e permitir exploração de código.

  5. execute_sql_query(sql: str) -> str

    • Nota: Esta é a mesma ferramenta listada em "Capacidades Principais", mas é repetida aqui no contexto de um fluxo de trabalho típico de agente envolvendo exploração do UC seguida de consulta.
    • Descrição: Executa uma consulta SQL fornecida contra o SQL warehouse do Databricks e retorna os resultados formatados.
    • Quando usar: Quando você precisa executar consultas SQL específicas, como SELECT, SHOW ou outras instruções DQL.
    • Argumentos:
      • sql: A string completa da consulta SQL a ser executada.

Configuração

Requisitos do Sistema

  • Python 3.10+
  • Se você planeja instalar via uv, certifique-se de que está instalado

Instalação

  1. Instale as dependências necessárias:
pip install -r requirements.txt

Ou se estiver usando uv:

uv pip install -r requirements.txt

  1. Configure suas variáveis de ambiente:

    Opção 1: Usando um arquivo .env (recomendado)

    Crie um arquivo .env no diretório raiz deste projeto com suas credenciais do Databricks:

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

    Opção 2: Definindo variáveis de ambiente diretamente

    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"

    Você pode encontrar seu ID do SQL Warehouse na interface do Databricks em "SQL Warehouses". O DATABRICKS_SQL_WAREHOUSE_ID é usado principalmente para buscar linhagem de tabelas e executar consultas SQL por meio da ferramenta execute_sql_query. As ferramentas de navegação de metadados (listar/descrever catálogos, esquemas, tabelas) usam as APIs gerais do UC do SDK do Databricks e não exigem estritamente um ID do SQL Warehouse, a menos que a linhagem seja solicitada.

Requisitos de Permissões

Antes de usar este servidor MCP, certifique-se de que a identidade associada ao DATABRICKS_TOKEN (por exemplo, um usuário ou service principal) tenha as permissões necessárias:

  1. Permissões do Unity Catalog:
    • USE CATALOG nos catálogos a serem acessados.
    • USE SCHEMA nos esquemas a serem acessados.
    • SELECT nas tabelas a serem consultadas ou descritas em detalhes (incluindo informações de colunas).
    • Para listar todos os catálogos, permissões apropriadas em nível de metastore podem ser necessárias, ou ele listará catálogos onde o usuário tem pelo menos USE CATALOG.
  2. Permissões do SQL Warehouse (para execute_sql_query e busca de linhagem):
    • Permissão CAN_USE no SQL Warehouse especificado por DATABRICKS_SQL_WAREHOUSE_ID.
  3. Permissões de Token:
    • O token de acesso pessoal ou token do service principal deve ter os escopos mínimos necessários. Para operações do Unity Catalog, isso normalmente envolve acesso ao workspace. Para execução de SQL, envolve permissões de SQL.
    • É fortemente recomendado usar um service principal com permissões estritamente definidas para cenários de produção ou automatizados.

Para práticas recomendadas de segurança, considere rotacionar regularmente seus tokens de acesso e auditar o histórico de consultas e os logs de auditoria do UC para monitorar o uso.

Executando o Servidor

Modo Autônomo

Para executar o servidor em modo autônomo (por exemplo, para testes com o Agent Composer):

python main.py

Isto iniciará o servidor MCP usando transporte stdio, que pode ser usado com o Agent Composer ou outros clientes MCP.

Usando com Cursor

Para usar este servidor MCP com Cursor, configure-o nas configurações do Cursor (~/.cursor/mcp.json):

  1. Crie um diretório .cursor no seu diretório inicial, se ele ainda não existir
  2. Crie ou edite o arquivo mcp.json nesse diretório:
mkdir -p ~/.cursor
touch ~/.cursor/mcp.json
  1. Adicione a seguinte configuração ao arquivo mcp.json, substituindo o caminho do diretório pelo caminho real onde você instalou este servidor:
{
    "mcpServers": {
        "databricks": {
            "command": "uv",
            "args": [
                "--directory",
                "/path/to/your/mcp-databricks-server",
                "run",
                "main.py"
            ]
        }
    }
}

Exemplo usando python:

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

Reinicie o Cursor para aplicar as alterações. Você pode então usar o agente databricks no Cursor.

Exemplo de Fluxo de Trabalho (para um Agente LLM)

Este servidor MCP capacita um agente LLM a navegar autonomamente pelo seu ambiente Databricks. A captura de tela a seguir ilustra uma interação típica em que o agente explora iterativamente schemas e tabelas, adaptando sua abordagem mesmo quando as consultas iniciais não retornam resultados, até conseguir recuperar com sucesso os dados solicitados.

Agent actively using MCP tools to find data Fig 3: Um agente LLM usando as ferramentas MCP do Databricks, demonstrando exploração iterativa e refinamento de consultas para localizar dados específicos de visualizações de página.

Um agente pode seguir este tipo de fluxo de trabalho:

  1. Descobrir catálogos disponíveis: list_uc_catalogs()
    • O agente decide que prod_catalog é relevante na lista.
  2. Explorar um catálogo específico: describe_uc_catalog(catalog_name="prod_catalog")
    • O agente vê sales_schema e inventory_schema.
  3. Explorar um schema específico (visualização rápida): describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema")
    • O agente vê nomes de tabelas como orders, customers.
  4. Obter estrutura detalhada da tabela (incluindo colunas para construção de consultas): describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema", include_columns=True)
    • Alternativamente, se uma tabela específica for de interesse: describe_uc_table(full_table_name="prod_catalog.sales_schema.orders")
  5. Analisar linhagem de dados e descobrir código de processamento: describe_uc_table(full_table_name="prod_catalog.sales_schema.orders", include_lineage=True)
    • O agente descobre tabelas upstream, dependências downstream e notebooks que processam esses dados
    • Por exemplo, vê que /Repos/production/etl/sales_processing.py grava nesta tabela
  6. Examinar a lógica de transformação de dados: O agente lê diretamente o arquivo do notebook /Repos/production/etl/sales_processing.py dentro do IDE/repositório
    • O agente analisa o código Python/SQL real para entender regras de negócio, verificações de qualidade de dados e lógica de transformação
  7. Construir e executar uma consulta: 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")

Gerenciando Metadados como Código com Terraform

Embora inserir metadados manualmente pela interface do Databricks seja uma opção, uma abordagem mais robusta e escalável é definir seus metadados do Unity Catalog como código. Ferramentas como Terraform permitem gerenciar declarativamente seus objetos de governança de dados, incluindo catálogos e schemas. Isso traz várias vantagens:

  • Controle de Versão: Suas definições de metadados podem ser armazenadas no Git, rastreadas e versionadas junto com seu outro código de infraestrutura.
  • Repetibilidade e Consistência: Garanta metadados consistentes entre ambientes (dev, staging, prod).
  • Automação: Integre o gerenciamento de metadados aos seus pipelines de CI/CD.
  • Manutenção Mais Fácil para Ativos Principais: Embora definir cada nova tabela como código possa ser complexo devido à sua natureza dinâmica, ativos principais como catálogos e schemas costumam ser mais estáveis e se beneficiam significativamente dessa abordagem. Manter suas definições e comentários como código garante uma base durável e bem documentada para o seu cenário de dados.

Aqui está um exemplo de como você pode definir um catálogo e seus schemas usando o provedor Databricks para Terraform:

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

Não se preocupe se você já possui catálogos e schemas existentes no Unity Catalog. Você não precisa recriá-los para gerenciar seus metadados como código. O Terraform fornece o comando terraform import, que permite trazer infraestrutura existente (incluindo ativos do Unity Catalog) para seu gerenciamento. Uma vez importado, você pode definir o recurso na sua configuração do Terraform e atualizar seletivamente atributos como o campo comment sem afetar o ativo em si. Por exemplo, após importar um schema existente, você pode adicionar ou atualizar seu comment no arquivo .tf, e o terraform apply aplicaria apenas essa alteração.

Adotar uma estratégia de metadados como código, especialmente para elementos fundamentais como catálogos e schemas, melhora muito a qualidade e a confiabilidade dos metadados que este servidor MCP utiliza. Isso, por sua vez, aumenta ainda mais a eficácia dos agentes de IA que interagem com seus dados do Databricks.

Para mais detalhes sobre o uso do Terraform com o Databricks Unity Catalog, consulte a documentação oficial:

Lidando com Consultas de Longa Duração

A ferramenta execute_sql_query utiliza o método execute_statement do SDK do Databricks. O parâmetro wait_timeout na função subjacente databricks_sdk_utils.execute_databricks_sql está definido como '50s'. Se uma consulta demorar mais que isso, o SDK pode retornar um ID de declaração para polling, mas a implementação atual da ferramenta efetivamente aguarda até essa duração por uma resposta síncrona. Para consultas muito longas, esse tempo limite pode ser atingido.

Dependências

  • databricks-sdk: Para interagir com as APIs REST do Databricks e o Unity Catalog.
  • python-dotenv: Para carregar variáveis de ambiente de um arquivo .env.
  • mcp[cli]: A biblioteca Model Context Protocol.
  • asyncio: Para operações assíncronas dentro do servidor MCP.
  • httpx (normalmente uma sub-dependência de databricks-sdk ou mcp): Para fazer requisições HTTP.