Cách Sử Dụng Máy Chủ Databricks MCP

💡

Trước khi tìm hiểu về Databricks MCP Server, hãy xem qua Apidog—một công cụ miễn phí được thiết kế để đơn giản hóa việc kiểm tra và tích hợp API. Với giao diện trực quan của Apidog, bạn có thể dễ dàng gỡ lỗi và tối ưu hóa quy trình làm việc với API của mình, làm cho quá trình phát triển thuận tiện hơn và tiết kiệm thời gian quý giá. Dù bạn đang phát triển API hay xử lý sự cố, Apidog có mọi thứ bạn cần để nâng cao quy trình làm việc của mình.

button

Model Context Protocol (MCP) là một giao thức mở được phát triển ban đầu bởi Anthropic, chuẩn hóa cách mà các ứng dụng cung cấp ngữ cảnh cho các Mô hình Ngôn ngữ Lớn (LLMs). Databricks MCP Server hoạt động như một cầu nối giữa nền tảng phân tích dữ liệu mạnh mẽ của Databricks và các mô hình AI, cho phép LLM tương tác với môi trường Databricks của bạn một cách lập trình. Tích hợp này cho phép trợ lý AI thực hiện các truy vấn SQL, liệt kê và lấy chi tiết về các thực thi công việc, và truy cập dữ liệu trong tài khoản Databricks của bạn mà không cần can thiệp trực tiếp từ con người.

Trong hướng dẫn toàn diện này, chúng ta sẽ đi qua quy trình cài đặt, cấu hình và sử dụng Databricks MCP Server để tăng cường các hoạt động và phân tích dữ liệu do AI dẫn dắt.

Hiểu về Model Context Protocol (MCP)

Trước khi đi vào triển khai cụ thể của Databricks, điều quan trọng là phải hiểu MCP là gì và tại sao nó lại quan trọng:

MCP thực hiện kiến trúc client/server với giao tiếp hai chiều:

server chứa các công cụ mà LLM có thể sử dụng
client sắp xếp các cuộc trò chuyện giữa người dùng, LLM và các công cụ

Các lợi ích chính khiến MCP có giá trị bao gồm:

Sự phân tách mối quan tâm: Các nhà phát triển công cụ có thể tạo ra và cập nhật khả năng một cách độc lập
Modularity: Các công cụ mới có thể được thêm vào một cách dần dần mà không cần thay đổi toàn bộ hệ thống
Tiêu chuẩn hóa: Một giao diện đồng nhất có nghĩa là ít công việc tích hợp hơn giữa các thành phần
Khả năng mở rộng: Các thành phần có thể được mở rộng độc lập dựa trên nhu cầu

Yêu cầu để chạy Databricks MCP Server

Trước khi thiết lập Databricks MCP Server, hãy đảm bảo bạn có:

Một tài khoản Databricks với quyền hạn phù hợp
Cài đặt Python 3.8+ trên hệ thống của bạn
Databricks CLI được cấu hình với xác thực
Hiểu biết cơ bản về các khái niệm không gian làm việc của Databricks
Quen thuộc với SQL và các công việc của Databricks

Bước 1: Thiết lập môi trường

Đầu tiên, hãy thiết lập môi trường của chúng ta bằng cách cài đặt các gói cần thiết:

# Tạo một môi trường ảo
python -m venv databricks-mcp-env
source databricks-mcp-env/bin/activate  # Trên Windows: databricks-mcp-env\\\\Scripts\\\\activate

# Cài đặt các gói yêu cầu
pip install databricks-sdk mcp-server httpx fastapi uvicorn

Bước 2: Tạo Databricks MCP Server

Tạo một tệp mới có tên databricks_mcp.py với cấu trúc mã như sau:

from typing import Any, Dict, List, Optional
from mcp.server.fastmcp import FastMCP
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.sql import ExecuteStatementRequest, StatementState

# Khởi tạo máy chủ FastMCP
mcp = FastMCP("databricks")

# Khởi tạo client Databricks
workspace = WorkspaceClient()

Bước 3: Triển khai các công cụ MCP cơ bản cho Databricks

Giờ đây, hãy triển khai các công cụ cốt lõi cho phép LLM tương tác với môi trường Databricks của bạn:

@mcp.tool()
async def run_sql_query(query: str, warehouse_id: Optional[str] = None) -> Dict[str, Any]:
    """Chạy một truy vấn SQL trong kho dữ liệu Databricks.

    Args:
        query: Truy vấn SQL để thực hiện
        warehouse_id: ID tùy chọn của kho dữ liệu SQL để sử dụng

    Returns:
        Từ điển chứa kết quả truy vấn
    """
    # Xác thực truy vấn để bảo mật
    if not query or any(keyword in query.lower() for keyword in ["drop", "delete", "truncate", "alter"]):
        return {"error": "Các thao tác SQL có thể phá hoại không được phép"}

    try:
        # Thực hiện truy vấn
        statement = workspace.statement_execution.execute_statement(
            warehouse_id=warehouse_id,
            statement=query,
            wait_timeout=60
        )

        # Xử lý kết quả
        if statement.status.state == StatementState.SUCCEEDED:
            # Chuyển đổi kết quả sang định dạng dễ sử dụng hơn
            columns = [col.name for col in statement.manifest.schema.columns]
            rows = []
            for chunk in workspace.statement_execution.get_statement_result_chunks(statement.statement_id):
                for row_data in chunk.data:
                    rows.append(dict(zip(columns, row_data)))

            return {
                "success": True,
                "columns": columns,
                "rows": rows,
                "row_count": len(rows)
            }
        else:
            return {
                "success": False,
                "state": statement.status.state,
                "error": statement.status.error
            }
    except Exception as e:
        return {"success": False, "error": str(e)}

Tiếp theo, hãy thêm một công cụ để liệt kê các cơ sở dữ liệu có sẵn:

@mcp.tool()
async def list_databases() -> List[str]:
    """Liệt kê tất cả các cơ sở dữ liệu có sẵn trong không gian làm việc Databricks.

    Returns:
        Danh sách tên cơ sở dữ liệu
    """
    try:
        # Chạy một truy vấn để lấy tất cả các cơ sở dữ liệu
        result = await run_sql_query("SHOW DATABASES")
        if result.get("success"):
            # Trích xuất tên cơ sở dữ liệu từ kết quả
            return [row.get("databaseName") for row in result.get("rows", [])]
        else:
            return [f"Lỗi khi liệt kê các cơ sở dữ liệu: {result.get('error')}"]
    except Exception as e:
        return [f"Lỗi: {str(e)}"]

Giờ đây, hãy triển khai một công cụ để lấy sơ đồ bảng:

@mcp.tool()
async def get_table_schema(database_name: str, table_name: str) -> Dict[str, Any]:
    """Lấy sơ đồ của một bảng cụ thể.

    Args:
        database_name: Tên của cơ sở dữ liệu
        table_name: Tên của bảng

    Returns:
        Từ điển chứa thông tin sơ đồ bảng
    """
    try:
        # Chạy một truy vấn để lấy sơ đồ bảng
        result = await run_sql_query(f"DESCRIBE {database_name}.{table_name}")
        if result.get("success"):
            return {
                "database": database_name,
                "table": table_name,
                "columns": result.get("rows", [])
            }
        else:
            return {"error": result.get("error")}
    except Exception as e:
        return {"error": str(e)}

Bước 4: Thêm các công cụ quản lý công việc

Hãy thêm các công cụ để quản lý và truy vấn các công việc Databricks:

@mcp.tool()
async def list_jobs(limit: int = 100) -> List[Dict[str, Any]]:
    """Liệt kê các công việc trong không gian làm việc Databricks.

    Args:
        limit: Số lượng tối đa các công việc để trả về

    Returns:
        Danh sách chi tiết công việc
    """
    try:
        jobs = workspace.jobs.list(limit=limit)
        result = []
        for job in jobs:
            result.append({
                "job_id": job.job_id,
                "name": job.settings.name,
                "created_by": job.created_by,
                "created_time": job.created_time
            })
        return result
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def get_job_runs(job_id: int, limit: int = 10) -> Dict[str, Any]:
    """Lấy các lần chạy gần đây cho một công việc cụ thể.

    Args:
        job_id: ID của công việc
        limit: Số lượng tối đa các lần chạy để trả về

    Returns:
        Từ điển chứa chi tiết về các lần chạy công việc
    """
    try:
        runs = workspace.jobs.list_runs(job_id=job_id, limit=limit)

        result = {
            "job_id": job_id,
            "runs": []
        }

        for run in runs.runs:
            run_info = {
                "run_id": run.run_id,
                "state": run.state.life_cycle_state,
                "result_state": run.state.result_state,
                "start_time": run.start_time,
                "setup_duration": run.setup_duration,
                "execution_duration": run.execution_duration,
                "cleanup_duration": run.cleanup_duration
            }
            result["runs"].append(run_info)

        return result
    except Exception as e:
        return {"error": str(e)}

Bước 5: Triển khai các công cụ quản lý cụm

Giờ đây hãy thêm các công cụ để tương tác với các cụm Databricks:

@mcp.tool()
async def list_clusters() -> List[Dict[str, Any]]:
    """Liệt kê tất cả các cụm trong không gian làm việc Databricks.

    Returns:
        Danh sách chi tiết cụm
    """
    try:
        clusters = workspace.clusters.list()
        result = []

        for cluster in clusters:
            cluster_info = {
                "cluster_id": cluster.cluster_id,
                "cluster_name": cluster.cluster_name,
                "state": cluster.state,
                "creator_username": cluster.creator_user_name,
                "spark_version": cluster.spark_version,
                "node_type_id": cluster.node_type_id
            }
            result.append(cluster_info)

        return result
    except Exception as e:
        return [{"error": str(e)}]

Bước 6: Chạy MCP Server

Cuối cùng, hãy thêm mã để chạy máy chủ MCP:

if __name__ == "__main__":
    # Chạy máy chủ với giao thức stdio cho tính tương thích với Claude Desktop
    # hoặc HTTP cho các khách hàng dựa trên web
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "--http":
        # Chạy dưới dạng máy chủ HTTP
        mcp.run(transport='http', host="127.0.0.1", port=8000)
    else:
        # Mặc định là stdio cho Claude Desktop
        mcp.run(transport='stdio')

Bước 7: Cấu hình MCP Client

Để sử dụng Databricks MCP Server với một client MCP như Claude Desktop, bạn sẽ cần tạo một cấu hình:

{
  "mcpServers": {
    "databricks": {
      "command": "python",
      "args": ["databricks_mcp.py"]
    }
  }
}

Nếu bạn muốn sử dụng giao thức HTTP, bạn có thể chạy máy chủ trong một quá trình riêng biệt và cấu hình client cho phù hợp:

{
  "mcpServers": {
    "databricks": {
      "url": "<http://localhost:8000>"
    }
  }
}

Bước 8: Xây dựng ứng dụng Client (Tùy chọn)

Để có trải nghiệm tùy chỉnh hơn, bạn có thể xây dựng một ứng dụng client khai thác Databricks MCP Server. Dưới đây là một ví dụ về một client dựa trên Python kết nối tới cả máy chủ MCP và mô hình Meta Llama của Databricks:

import asyncio
from contextlib import AsyncExitStack
from typing import Optional, Dict, Any

from databricks.sdk import WorkspaceClient
from mcp.client import ClientSession, stdio_client
from mcp.stdio import StdioServerParameters

class DatabricksMCPClient:
    def __init__(self, llm_endpoint_name: str = "databricks-llama3-70b-instruct"):
        # Khởi tạo các đối tượng phiên và client
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()

        # Khởi tạo client Databricks
        self.workspace = WorkspaceClient()
        self.llm_endpoint_name = llm_endpoint_name
        self.openai_client = self.workspace.serving_endpoints.get_open_ai_client()
        print(f"Đã khởi tạo client tương thích với OpenAI cho {llm_endpoint_name}")

    async def connect_to_server(self, server_script_path: str):
        """Kết nối đến máy chủ Databricks MCP"""
        server_params = StdioServerParameters(
            command="python",
            args=[server_script_path],
            env=None
        )
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))

        await self.session.initialize()

        # Liệt kê các công cụ có sẵn
        response = await self.session.list_tools()
        tools = response.tools
        print("\\\\nKết nối đến máy chủ Databricks MCP với các công cụ:", [tool.name for tool in tools])

        # Hiển thị các mô tả công cụ
        print("\\\\nCác Công cụ Databricks Có sẵn:")
        for tool in tools:
            print(f"- {tool.name}: {tool.description}")

    # Thêm các phương thức bổ sung để xử lý truy vấn, v.v.

async def main():
    client = DatabricksMCPClient()
    try:
        await client.connect_to_server("databricks_mcp.py")
        # Thêm vòng lặp tương tác hoặc chức năng khác
    finally:
        await client.exit_stack.aclose()

if __name__ == "__main__":
    asyncio.run(main())

Các thực hành tốt nhất để sử dụng Databricks MCP Server

Bảo mật làưu tiên hàng đầu: Luôn áp dụng kiểm soát truy cập và xác thực truy vấn SQL để ngăn chặn các thao tác phá hoại.
Xử lý lỗi: Triển khai xử lý lỗi mạnh mẽ trong cả ứng dụng máy chủ và client.
Giới hạn tần suất: Thêm giới hạn tần suất để ngăn chặn các cuộc gọi API quá mức tới Databricks.
Ghi nhật ký: Triển khai ghi nhật ký toàn diện để phục vụ việc gỡ lỗi và kiểm toán.
Phân trang: Đối với các tập kết quả lớn, triển khai phân trang để tránh các vấn đề về bộ nhớ.
Bảo mật thông tin xác thực: Sử dụng biến môi trường hoặc lưu trữ thông tin xác thực an toàn thay vì mã hóa trực tiếp thông tin xác thực.
Tài liệu công cụ: Viết mô tả rõ ràng, chi tiết cho từng công cụ để giúp LLM sử dụng chúng một cách chính xác.

Tính năng nâng cao

Thêm hỗ trợ cho Delta Live Tables

@mcp.tool()
async def list_delta_pipelines() -> List[Dict[str, Any]]:
    """Liệt kê tất cả các pipeline Delta Live Table trong không gian làm việc.

    Returns:
        Danh sách chi tiết pipeline
    """
    try:
        pipelines = workspace.pipelines.list()
        result = []

        for pipeline in pipelines:
            pipeline_info = {
                "pipeline_id": pipeline.pipeline_id,
                "name": pipeline.name,
                "state": pipeline.state,
                "creator": pipeline.creator,
                "target": pipeline.target
            }
            result.append(pipeline_info)

        return result
    except Exception as e:
        return [{"error": str(e)}]

Thêm quản lý Sổ tay

@mcp.tool()
async def list_notebooks(path: str = "/") -> List[Dict[str, Any]]:
    """Liệt kê các sổ tay trong đường dẫn đã chỉ định trong không gian làm việc.

    Args:
        path: Đường dẫn không gian làm việc để liệt kê các sổ tay từ đó

    Returns:
        Danh sách chi tiết sổ tay
    """
    try:
        objects = workspace.workspace.list(path)
        notebooks = []

        for obj in objects:
            if obj.object_type == "NOTEBOOK":
                notebooks.append({
                    "path": obj.path,
                    "name": obj.path.split("/")[-1],
                    "language": obj.language
                })

        return notebooks
    except Exception as e:
        return [{"error": str(e)}]

Kết luận

Databricks MCP Server cung cấp một cách mạnh mẽ để kết nối giữa các mô hình AI và cơ sở hạ tầng dữ liệu của tổ chức bạn. Bằng việc theo dõi hướng dẫn toàn diện này, bạn đã học cách:

Cài đặt và cấu hình Databricks MCP Server
Triển khai các công cụ cốt lõi cho việc thực hiện truy vấn SQL và khám phá cơ sở dữ liệu
Thêm khả năng quản lý công việc và cụm
Cấu hình các client MCP để tương tác với máy chủ của bạn
Tùy chọn xây dựng một ứng dụng client tùy chỉnh

Sự tích hợp này giúp các trợ lý AI của bạn có thể trực tiếp khai thác khả năng xử lý dữ liệu mạnh mẽ của Databricks, cho phép phân tích dữ liệu, giám sát và tự động hóa tinh vi hơn. Khi hệ sinh thái Model Context Protocol tiếp tục phát triển, máy chủ Databricks MCP của bạn có thể được mở rộng với các công cụ và khả năng bổ sung để đáp ứng nhu cầu cụ thể của tổ chức bạn.

Bằng việc chấp nhận sự tích hợp này, bạn đang thực hiện một bước tiến lớn hướng tới việc tạo ra các hệ thống AI thông minh, nhận thức về dữ liệu có thể cung cấp những cái nhìn sâu sắc hơn và hỗ trợ hiệu quả hơn cho quy trình phân tích dữ liệu của bạn.