Databricks Genie MCP

🚀 Databricks Genie API MCP Server

This project implements a Model Context Protocol (MCP) server that exposes the Databricks Genie API as a set of tools. These tools can be integrated into IDEs or other applications to enhance the user experience.

✨ Features

• Start Conversation: Initiate a new conversation in the Genie space.
• Create Message: Add a new message to an existing conversation.
• Get Message: Retrieve a message from a conversation.
• Get Message Attachment Query Result: Obtain the SQL query results from a message attachment.
• Execute Message Attachment Query: Execute the SQL query in a message attachment.
• Get Space Information: Fetch details about the Genie space.
• Generate Download Full Query Result: Initiate the download of full query results.
• Poll Message Until Complete: Continuously check a message until it reaches a terminal state.

📦 Installation

Prerequisites

1. Install Python 3.8 or higher.
2. Install the Databricks CLI:

   pip install databricks-cli
   

3. Configure the Databricks CLI:

   dbfs config set <your-workspace-url> <access-token>
   

Setup Steps

1. Clone the repository and install dependencies:

   git clone https://github.com/your-repository.git
   cd genie_mcp_server
   pip install -r requirements.txt
   

2. Configure environment variables or the .env file with the following settings:

   DATABRICKS_WORKSPACE_URL=your_workspace_url
   DATABRICKS_TOKEN=your_access_token
   

3. Start the server:

   python server.py
   

💻 Usage Examples

Integration with Claude Desktop

Configure the MCP server in claude_desktop_config.json:

{
  "mcpServers": {
    "databricks-genie": {
      "type": "stdio",
      "command": "/path/to/python",
      "args": ["server.py"],
      "env": {
        "DATABRICKS_WORKSPACE_URL": "your_workspace_url",
        "DATABRICKS_TOKEN": "your_access_token"
      }
    }
  }
}

Tool List

• start_conversation: Create a new conversation in the Genie space.
• create_message: Add new content to an existing conversation.
• get_message: Retrieve detailed information about a specific message.
• get_message_attachment_query_result: Obtain SQL query results from a message attachment.
• execute_message_attachment_query: Execute an SQL code snippet in the Genie space.
• get_space: Get detailed information about a specified space.
• generate_download_full_query_result: Initiate the download process for full query results.
• poll_message_until_complete: Continuously check the message status until it is either completed or failed.

📚 Documentation

Troubleshooting

• Authentication Issues: Check the Databricks credentials in the .env file and ensure that the user has appropriate permissions.
• Connection Issues (Claude Desktop):
  • Ensure that the paths in claude_desktop_config.json are correct.
  • Verify that the Python interpreter path exists and is executable.
  • Check the log files, such as ~/Library/Logs/Claude/ or %APPDATA%\Claude\logs, especially mcp.log and mcp-server-databricks-genie.log.
• SQL Execution Errors: Confirm that the service principal or user has execution permissions for the SQL warehouse and can access the relevant Unity Catalog data.

Security Considerations

• Credential Management: Never hard - code credentials. Use environment variables (.env file) or a secure credential store. Ensure that the .env file is not committed to the version control system.
• MCP Security: This server runs locally, using the user's permissions and Databricks credentials. Clients such as Claude Desktop must obtain user consent before executing tools (refer to the MCP security specifications).
• Production Deployment: Do not directly expose this local server for broader use. Collaborate with the security team to determine appropriate hosting settings, network controls, and possible MCP server - level authentication mechanisms.

🤝 Contribution

Contributions are welcome! Please participate in the project by submitting issues or pull requests.

📄 License

This software is licensed by Databricks, Inc. under specific terms. Please review the LICENSE file for the full license agreement.