Databricks MCP Server

Interactúa con los metadatos de Databricks Unity Catalog para explorar catálogos, esquemas, tablas y ejecutar consultas SQL.

GitHub
42
Prueba este MCPPatrocinado

Documentación

Listed on Spark Install via Spark

Databricks MCP Server

Motivación

Databricks Unity Catalog (UC) permite documentar detalladamente tus activos de datos, incluidos catálogos, esquemas, tablas y columnas. Documentar estos activos a fondo requiere una inversión de tiempo. Una pregunta común es: ¿cuáles son los beneficios prácticos de esta entrada de metadatos detallada?

Este servidor MCP proporciona una justificación sólida para ese esfuerzo. Permite que los modelos de lenguaje grandes (LLM) accedan directamente y utilicen estos metadatos de Unity Catalog. Cuanto más exhaustivamente estén descritos tus datos en UC, más eficazmente un agente LLM podrá comprender tu entorno de Databricks. Esta comprensión más profunda es crucial para que el agente construya de forma autónoma consultas SQL más inteligentes y precisas para satisfacer solicitudes de datos.

Descripción general

Este servidor de Protocolo de Contexto de Modelo (MCP) está diseñado para interactuar con Databricks, con un fuerte enfoque en aprovechar los metadatos de Unity Catalog (UC) y permitir una exploración integral de la linaje de datos. El objetivo principal es equipar a un agente de IA con un conjunto completo de herramientas, permitiéndole volverse independiente para responder preguntas sobre tus datos. Al explorar UC de forma autónoma, comprender las estructuras de datos, analizar el linaje de datos (incluidas las dependencias de notebooks y trabajos) y ejecutar consultas SQL, el agente puede cumplir solicitudes de datos sin intervención humana directa en cada paso.

Más allá de la navegación tradicional de catálogos, este servidor permite a los agentes descubrir y analizar el código real que procesa tus datos. Mediante capacidades mejoradas de linaje, los agentes pueden identificar notebooks y trabajos que leen o escriben en tablas, y luego examinar la lógica de transformación real, las reglas de negocio y los controles de calidad de datos implementados en esos notebooks. Esto crea un poderoso bucle de retroalimentación donde los agentes no solo entienden qué datos existen, sino también cómo se procesan y transforman.

Cuando se usa en un modo de agente, puede iterar con éxito sobre una serie de solicitudes para realizar tareas complejas, incluido el descubrimiento de datos, el análisis de impacto y la exploración de código.

Beneficios prácticos de los metadatos de UC para agentes de IA

Las herramientas proporcionadas por este servidor MCP están diseñadas para analizar y presentar las descripciones que has agregado a Unity Catalog, al mismo tiempo que permiten una exploración profunda de tu código de procesamiento de datos. Esto ofrece ventajas tangibles para los agentes basados en LLM, impactando directamente su capacidad para generar SQL útil y comprender tu ecosistema de datos:

  • Contexto de datos más claro: Los agentes pueden comprender rápidamente el propósito de tablas y columnas, reduciendo la ambigüedad. Esta comprensión fundamental es el primer paso hacia una formulación correcta de consultas.
  • Generación de consultas más precisa: El acceso a descripciones, tipos de datos y relaciones ayuda a los agentes a construir consultas SQL con mayor precisión y corrección semántica.
  • Exploración eficiente de datos para la planificación de consultas: Los metadatos permiten a los agentes navegar por catálogos y esquemas de manera más efectiva, permitiéndoles identificar las tablas y columnas correctas para incluir en sus consultas SQL.
  • Linaje de datos integral: Más allá de las relaciones tabla a tabla, los agentes pueden descubrir notebooks y trabajos que procesan datos, lo que permite el análisis de impacto y la depuración de problemas en los pipelines de datos.
  • Comprensión a nivel de código: A través de la exploración del contenido de los notebooks, los agentes pueden analizar la lógica de transformación real, las reglas de negocio y los controles de calidad de datos, proporcionando información más profunda sobre cómo se procesan y transforman los datos.
  • Análisis de flujo de datos de extremo a extremo: Los agentes pueden rastrear los datos desde la ingesta bruta a través de los pipelines de transformación hasta el consumo final, comprendiendo tanto la estructura como la lógica de procesamiento en cada paso.

Los metadatos bien documentados en Unity Catalog, cuando se accede a ellos a través de este servidor, permiten que un agente LLM opere con mejor información y tome decisiones más informadas, culminando en la generación de consultas SQL más efectivas. Por ejemplo, las descripciones de esquemas ayudan al agente a identificar fuentes de datos relevantes para una consulta:

Schema Description in Unity Catalog Fig 1: Un esquema en Unity Catalog con descripciones proporcionadas por el usuario. Este servidor MCP hace que esta información sea directamente accesible para un LLM, informando su estrategia de consulta.

De manera similar, los comentarios detallados a nivel de columna aclaran la semántica de cada campo, lo cual es crucial para construir condiciones y selecciones SQL precisas:

Table Column Descriptions in Unity Catalog Fig 2: Descripciones a nivel de columna en Unity Catalog. Estos detalles se pasan al LLM, ayudando a su comprensión de la estructura de datos para la generación precisa de SQL.

Herramientas y funciones disponibles

Este servidor MCP proporciona un conjunto de herramientas diseñadas para empoderar a un agente LLM que interactúa con Databricks:

Capacidades principales:

  • Ejecutar consultas SQL: Ejecuta consultas SQL arbitrarias usando el SDK de Databricks a través de la herramienta execute_sql_query(sql: str). Esto es ideal para la recuperación de datos específica u operaciones complejas.
  • Salida enfocada en LLM: Todas las herramientas descriptivas devuelven información en formato Markdown, optimizado para el consumo por modelos de lenguaje grandes, facilitando que los agentes analicen y comprendan el contexto.

Herramientas de exploración de Unity Catalog:

El servidor proporciona las siguientes herramientas para navegar y comprender tus activos de Unity Catalog. Están diseñadas para ser utilizadas por un agente LLM para recopilar contexto antes de construir consultas o tomar decisiones, de manera agéntica.

  1. list_uc_catalogs() -> str

    • Descripción: Lista todos los catálogos de Unity disponibles con sus nombres, descripciones y tipos.
    • Cuándo usarla: Como punto de partida para descubrir fuentes de datos disponibles cuando no conoces nombres de catálogos específicos. Proporciona una visión general de alto nivel de todos los catálogos accesibles en el espacio de trabajo.

  2. describe_uc_catalog(catalog_name: str) -> str

    • Descripción: Proporciona un resumen de un catálogo de Unity específico, listando todos sus esquemas con sus nombres y descripciones.
    • Cuándo usarla: Cuando conoces el nombre del catálogo y necesitas descubrir los esquemas dentro de él. Esto suele ser un precursor para describir un esquema o tabla específica.
    • Argumentos:
      • catalog_name: El nombre del catálogo de Unity a describir (por ejemplo, prod, dev, system).

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

    • Descripción: Proporciona información detallada sobre un esquema específico dentro de un catálogo de Unity. Devuelve todas las tablas en el esquema, opcionalmente incluyendo sus detalles de columnas.
    • Cuándo usarla: Para comprender el contenido de un esquema, principalmente sus tablas. Establece include_columns=True para obtener información de columnas, crucial para la construcción de consultas pero hace la salida más larga. Si include_columns=False, solo se muestran nombres y descripciones de tablas, útil para una visión general más rápida.
    • Argumentos:
      • catalog_name: El nombre del catálogo que contiene el esquema.
      • schema_name: El nombre del esquema a describir.
      • include_columns: Si es True, lista las tablas con sus columnas. Por defecto es False para un resumen más breve.

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

    • Descripción: Proporciona una descripción detallada de una tabla específica de Unity Catalog con capacidades integrales de linaje.
    • Cuándo usarla: Para comprender la estructura (columnas, tipos de datos, particionamiento) de una sola tabla. Esto es esencial antes de construir consultas SQL contra la tabla. Opcionalmente, puede incluir información de linaje integral que va más allá de las dependencias tradicionales tabla a tabla:
      • Linaje de tablas: Tablas ascendentes (tablas de las que esta tabla lee) y tablas descendentes (tablas que leen de esta tabla)
      • Linaje de notebooks y trabajos: Notebooks que leen o escriben en esta tabla, incluido el nombre del notebook, la ruta del espacio de trabajo, la información del trabajo de Databricks asociado (nombre del trabajo, ID, detalles de la tarea)
      • Descubrimiento de código: El linaje proporciona rutas de notebooks que permiten a un agente leer directamente los archivos de notebook dentro del repositorio/espacio de trabajo actual, permitiendo el análisis de la lógica real de transformación de datos
    • Argumentos:
      • full_table_name: El nombre completamente calificado de tres partes de la tabla (por ejemplo, catalog.schema.table).
      • include_lineage: Establecer a True para obtener linaje integral (tablas, notebooks, trabajos). Por defecto es False. Puede tardar más en recuperarse pero proporciona un contexto rico para comprender las dependencias de datos y permitir la exploración de código.

  5. execute_sql_query(sql: str) -> str

    • Nota: Esta es la misma herramienta listada bajo "Capacidades principales" pero se repite aquí en el contexto de un flujo de trabajo típico de agente que involucra exploración de UC seguida de consultas.
    • Descripción: Ejecuta una consulta SQL dada contra el almacén de SQL de Databricks y devuelve los resultados formateados.
    • Cuándo usarla: Cuando necesitas ejecutar consultas SQL específicas, como SELECT, SHOW u otras declaraciones DQL.
    • Argumentos:
      • sql: La cadena de consulta SQL completa a ejecutar.

Configuración

Requisitos del sistema

  • Python 3.10+
  • Si planeas instalar a través de uv, asegúrate de que esté instalado

Instalación

  1. Instala las dependencias requeridas:
pip install -r requirements.txt

O si usas uv:

uv pip install -r requirements.txt

  1. Configura tus variables de entorno:

    Opción 1: Usar un archivo .env (recomendado)

    Crea un archivo .env en el directorio raíz de este proyecto con tus credenciales de Databricks:

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

    Opción 2: Configurar variables de entorno directamente

    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"

    Puedes encontrar tu ID de almacén de SQL en la interfaz de Databricks bajo "SQL Warehouses". El DATABRICKS_SQL_WAREHOUSE_ID se utiliza principalmente para obtener el linaje de tablas y ejecutar consultas SQL a través de la herramienta execute_sql_query. Las herramientas de navegación de metadatos (listar/describir catálogos, esquemas, tablas) utilizan las API generales de UC del SDK de Databricks y no requieren estrictamente un ID de almacén de SQL a menos que se solicite linaje.

Requisitos de permisos

Antes de usar este servidor MCP, asegúrate de que la identidad asociada con el DATABRICKS_TOKEN (por ejemplo, un usuario o principal de servicio) tenga los permisos necesarios:

  1. Permisos de Unity Catalog:
    • USE CATALOG en los catálogos a los que se accederá.
    • USE SCHEMA en los esquemas a los que se accederá.
    • SELECT en las tablas que se consultarán o describirán en detalle (incluida la información de columnas).
    • Para listar todos los catálogos, podrían necesitarse permisos a nivel de metastore o se listarán los catálogos donde el usuario tenga al menos USE CATALOG.
  2. Permisos de almacén de SQL (para execute_sql_query y obtención de linaje):
    • CAN_USE permiso en el almacén de SQL especificado por DATABRICKS_SQL_WAREHOUSE_ID.
  3. Permisos de token:
    • El token de acceso personal o el token de principal de servicio debe tener los alcances mínimos necesarios. Para operaciones de Unity Catalog, esto típicamente involucra acceso al espacio de trabajo. Para la ejecución de SQL, involucra permisos de SQL.
    • Se recomienda encarecidamente usar un principal de servicio con permisos definidos de manera estrecha para escenarios de producción o automatizados.

Para las mejores prácticas de seguridad, considera rotar regularmente tus tokens de acceso y auditar el historial de consultas y los registros de auditoría de UC para monitorear el uso.

Ejecución del servidor

Modo independiente

Para ejecutar el servidor en modo independiente (por ejemplo, para pruebas con Agent Composer):

python main.py

Esto iniciará el servidor MCP usando transporte stdio, que puede usarse con Agent Composer u otros clientes MCP.

Uso con Cursor

Para usar este servidor MCP con Cursor, configúralo en la configuración de Cursor (~/.cursor/mcp.json):

  1. Crea un directorio .cursor en tu directorio de inicio si aún no existe
  2. Crea o edita el archivo mcp.json en ese directorio:
mkdir -p ~/.cursor
touch ~/.cursor/mcp.json
  1. Agrega la siguiente configuración al archivo mcp.json, reemplazando la ruta del directorio con la ruta real donde has instalado este servidor:
{
    "mcpServers": {
        "databricks": {
            "command": "uv",
            "args": [
                "--directory",
                "/path/to/your/mcp-databricks-server",
                "run",
                "main.py"
            ]
        }
    }
}

Ejemplo usando python:

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

Reinicia Cursor para aplicar los cambios. Luego puedes usar el agente databricks en Cursor.

Flujo de trabajo de uso de ejemplo (para un agente LLM)

Este servidor MCP permite que un agente LLM navegue de forma autónoma por tu entorno de Databricks. La siguiente captura de pantalla ilustra una interacción típica donde el agente explora iterativamente esquemas y tablas, adaptando su enfoque incluso cuando las consultas iniciales no producen resultados, hasta que recupera exitosamente los datos solicitados.

Agent actively using MCP tools to find data Fig 3: Un agente LLM usando las herramientas MCP de Databricks, demostrando exploración iterativa y refinamiento de consultas para localizar datos específicos de vistas de página.

Un agente podría seguir este tipo de flujo de trabajo:

  1. Descubrir catálogos disponibles: list_uc_catalogs()
    • El agente decide que prod_catalog es relevante de la lista.
  2. Explorar un catálogo específico: describe_uc_catalog(catalog_name="prod_catalog")
    • El agente ve sales_schema y inventory_schema.
  3. Explorar un esquema específico (vista rápida): describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema")
    • El agente ve nombres de tablas como orders, customers.
  4. Obtener estructura detallada de la tabla (incluyendo columnas para construir consultas): describe_uc_schema(catalog_name="prod_catalog", schema_name="sales_schema", include_columns=True)
    • Alternativamente, si una tabla específica es de interés: describe_uc_table(full_table_name="prod_catalog.sales_schema.orders")
  5. Analizar linaje de datos y descubrir código de procesamiento: describe_uc_table(full_table_name="prod_catalog.sales_schema.orders", include_lineage=True)
    • El agente descubre tablas upstream, dependencias downstream y notebooks que procesan estos datos
    • Por ejemplo, ve que /Repos/production/etl/sales_processing.py escribe en esta tabla
  6. Examinar lógica de transformación de datos: El agente lee directamente el archivo de notebook /Repos/production/etl/sales_processing.py dentro del IDE/repositorio
    • El agente analiza el código real de Python/SQL para entender reglas de negocio, controles de calidad de datos y lógica de transformación
  7. Construir y ejecutar una 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")

Gestión de metadatos como código con Terraform

Si bien ingresar metadatos manualmente a través de la interfaz de Databricks es una opción, un enfoque más robusto y escalable es definir los metadatos de Unity Catalog como código. Herramientas como Terraform te permiten gestionar declarativamente tus objetos de gobernanza de datos, incluidos catálogos y esquemas. Esto trae varias ventajas:

  • Control de versiones: Tus definiciones de metadatos pueden almacenarse en Git, rastrearse y versionarse junto con tu otro código de infraestructura.
  • Repetibilidad y consistencia: Asegura metadatos consistentes en todos los entornos (dev, staging, prod).
  • Automatización: Integra la gestión de metadatos en tus pipelines de CI/CD.
  • Mantenimiento más fácil para activos principales: Si bien definir cada nueva tabla como código puede ser complejo debido a su naturaleza dinámica, los activos principales como catálogos y esquemas suelen ser más estables y se benefician significativamente de este enfoque. Mantener sus definiciones y comentarios como código asegura una base duradera y bien documentada para tu panorama de datos.

Aquí hay un ejemplo de cómo podrías definir un catálogo y sus esquemas usando el proveedor de 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."
}

No te preocupes si ya tienes catálogos y esquemas existentes en Unity Catalog. No necesitas recrearlos para gestionar sus metadatos como código. Terraform proporciona el comando terraform import, que te permite traer infraestructura existente (incluidos activos de Unity Catalog) bajo su gestión. Una vez importado, puedes definir el recurso en tu configuración de Terraform y actualizar selectivamente atributos como el campo comment sin afectar el activo en sí. Por ejemplo, después de importar un esquema existente, podrías agregar o actualizar su comment en tu archivo .tf, y terraform apply solo aplicaría ese cambio.

Adoptar una estrategia de metadatos como código, especialmente para elementos fundamentales como catálogos y esquemas, mejora enormemente la calidad y confiabilidad de los metadatos que este servidor MCP aprovecha. Esto, a su vez, mejora aún más la efectividad de los agentes de IA que interactúan con tus datos de Databricks.

Para más detalles sobre el uso de Terraform con Databricks Unity Catalog, consulta la documentación oficial:

Manejo de consultas de larga duración

La herramienta execute_sql_query utiliza el método execute_statement del SDK de Databricks. El parámetro wait_timeout en la función subyacente databricks_sdk_utils.execute_databricks_sql está establecido en '50s'. Si una consulta se ejecuta más tiempo que esto, el SDK puede devolver un ID de declaración para sondeo, pero la implementación actual de la herramienta espera efectivamente hasta esta duración para una respuesta similar a la síncrona. Para consultas de muy larga duración, este tiempo de espera podría alcanzarse.

Dependencias

  • databricks-sdk: Para interactuar con las API REST de Databricks y Unity Catalog.
  • python-dotenv: Para cargar variables de entorno desde un archivo .env.
  • mcp[cli]: La biblioteca del Protocolo de Contexto de Modelo.
  • asyncio: Para operaciones asíncronas dentro del servidor MCP.
  • httpx (típicamente una sub-dependencia de databricks-sdk o mcp): Para realizar solicitudes HTTP.