#!/usr/bin/env node
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
CallToolRequestSchema,
ListToolsRequestSchema,
Tool,
} from '@modelcontextprotocol/sdk/types.js';
import * as dotenv from 'dotenv';
import { DatabricksClient } from './databricks-client.js';
import { clusterTools } from './tools/clusters.js';
import { notebookTools } from './tools/notebooks.js';
import { jobTools } from './tools/jobs.js';
import { userTools } from './tools/users.js';
import { unityCatalogTools } from './tools/unity-catalog.js';
import { permissionTools } from './tools/permissions.js';
import { finopsTools } from './tools/finops.js';
// Load environment variables
dotenv.config();
// Validate required environment variables
const DATABRICKS_HOST = process.env.DATABRICKS_HOST;
const DATABRICKS_TOKEN = process.env.DATABRICKS_TOKEN;
if (!DATABRICKS_HOST) {
console.error('Error: DATABRICKS_HOST environment variable is required');
process.exit(1);
}
if (!DATABRICKS_TOKEN) {
console.error('Error: DATABRICKS_TOKEN environment variable is required');
process.exit(1);
}
// Initialize Databricks client
const databricksClient = new DatabricksClient({
host: DATABRICKS_HOST,
token: DATABRICKS_TOKEN,
});
// Combine all tools
const allTools = {
...clusterTools,
...notebookTools,
...jobTools,
...userTools,
...unityCatalogTools,
...permissionTools,
...finopsTools,
};
// Create MCP server
const server = new Server(
{
name: 'databricks-mcp-server',
version: '1.0.0',
},
{
capabilities: {
tools: {},
},
}
);
// Register tool list handler
server.setRequestHandler(ListToolsRequestSchema, async () => {
const tools: Tool[] = Object.entries(allTools).map(([name, tool]) => ({
name,
description: tool.description,
inputSchema: tool.inputSchema as any,
}));
return { tools };
});
// Register tool call handler
server.setRequestHandler(CallToolRequestSchema, async (request) => {
const toolName = request.params.name;
const tool = allTools[toolName as keyof typeof allTools];
if (!tool) {
throw new Error(`Unknown tool: ${toolName}`);
}
try {
const result = await tool.handler(databricksClient, request.params.arguments || {});
return result;
} catch (error: any) {
// Return error information in a structured format
return {
content: [
{
type: 'text',
text: JSON.stringify(
{
error: true,
message: error.message || 'An unknown error occurred',
tool: toolName,
details: error.response?.data || error.toString(),
},
null,
2
),
},
],
isError: true,
};
}
});
// Error handler
server.onerror = (error) => {
console.error('[MCP Error]', error);
};
// Graceful shutdown handler
process.on('SIGINT', async () => {
await server.close();
process.exit(0);
});
// Start the server
async function main() {
const transport = new StdioServerTransport();
await server.connect(transport);
console.error('Databricks MCP Server running on stdio');
}
main().catch((error) => {
console.error('Fatal error in main():', error);
process.exit(1);
});
