mcp_server.py

#!/usr/bin/env python3

import json
import sys
import logging
from typing import Any, Dict
import db_handler

def setup_logging():
    """Setup logging to stderr for MCP server."""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        stream=sys.stderr
    )
    return logging.getLogger(__name__)

logger = setup_logging()

class MCPServer:
    def __init__(self):
        self.initialized = False
        self.tools = [
            {
                "name": "connect_database",
                "description": "Connect to a database (MySQL or PostgreSQL) with provided credentials",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "connection_name": {
                            "type": "string",
                            "description": "Unique name for this database connection"
                        },
                        "db_type": {
                            "type": "string",
                            "description": "Database type: 'mysql' or 'postgresql'"
                        },
                        "host": {
                            "type": "string",
                            "description": "Database host"
                        },
                        "port": {
                            "type": "integer",
                            "description": "Database port"
                        },
                        "user": {
                            "type": "string",
                            "description": "Database username"
                        },
                        "password": {
                            "type": "string",
                            "description": "Database password"
                        },
                        "database": {
                            "type": "string",
                            "description": "Database name"
                        }
                    },
                    "required": ["connection_name", "db_type", "host", "port", "user", "password", "database"]
                }
            },
            {
                "name": "execute_query",
                "description": "Execute a READ-ONLY SQL query on a connected database.",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "connection_name": {
                            "type": "string",
                            "description": "Name of the database connection to use"
                        },
                        "query": {
                            "type": "string",
                            "description": "READ-ONLY SQL query to execute"
                        },
                        "parameters": {
                            "type": "array",
                            "description": "Optional parameters for parameterized queries",
                            "items": {}
                        }
                    },
                    "required": ["connection_name", "query"]
                }
            },
            {
                "name": "list_tables",
                "description": "List all tables in the connected database",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "connection_name": {
                            "type": "string",
                            "description": "Name of the database connection to use"
                        }
                    },
                    "required": ["connection_name"]
                }
            },
            {
                "name": "describe_table",
                "description": "Get the structure/schema of a specific table",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "connection_name": {
                            "type": "string",
                            "description": "Name of the database connection to use"
                        },
                        "table_name": {
                            "type": "string",
                            "description": "Name of the table to describe"
                        }
                    },
                    "required": ["connection_name", "table_name"]
                }
            }
        ]

    def handle_initialize(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """Handle the initialize request."""
        self.initialized = True
        logger.info("MCP Server initialized")
        return {
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": {}
            },
            "serverInfo": {
                "name": "sql-local-mcp",
                "version": "1.1.0"
            }
        }

    def handle_list_tools(self) -> Dict[str, Any]:
        """Return the list of available tools."""
        return {"tools": self.tools}

    def handle_call_tool(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """Handle tool calls from MCP client."""
        tool_name = params.get("name")
        arguments = params.get("arguments", {})
        
        logger.info(f"Tool call: {tool_name} with arguments: {arguments}")
        
        try:
            if tool_name == "connect_database":
                result = db_handler.connect_database(**arguments)
            elif tool_name == "execute_query":
                result = db_handler.execute_query(**arguments)
            elif tool_name == "list_tables":
                result = db_handler.list_tables(**arguments)
            elif tool_name == "describe_table":
                result = db_handler.describe_table(**arguments)
            else:
                result = {"success": False, "error": f"Unknown tool: {tool_name}"}
                
            return {
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps(result, indent=2, default=str)
                    }
                ]
            }
        except Exception as e:
            logger.error(f"Error handling tool call {tool_name}: {str(e)}")
            return {
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps({"success": False, "error": str(e)}, indent=2)
                    }
                ]
            }

    def send_response(self, response: Dict[str, Any]):
        """Send a response to stdout and flush."""
        output = json.dumps(response)
        print(output)
        sys.stdout.flush()
        logger.info(f"Sent: {response.get('method', 'response')}")

    def send_error_response(self, request_id: Any, code: int, message: str):
        """Send an error response."""
        error_response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": code,
                "message": message
            }
        }
        self.send_response(error_response)

    def run(self):
        """Main server loop."""
        logger.info("Starting SQL Local MCP Server...")
        
        try:
            while True:
                line = sys.stdin.readline()
                if not line:
                    logger.info("EOF reached, shutting down")
                    break

                try:
                    request = json.loads(line)
                    logger.info(f"Received: {request}")
                    
                    method = request.get("method")
                    params = request.get("params", {})
                    request_id = request.get("id")
                    
                    if request_id is None:
                        if method == "notifications/initialized":
                            logger.info("Client sent initialized notification")
                        continue

                    if method == "initialize":
                        result = self.handle_initialize(params)
                    elif not self.initialized:
                        self.send_error_response(request_id, -32002, "Server not initialized")
                        continue
                    elif method == "tools/list":
                        result = self.handle_list_tools()
                    elif method == "tools/call":
                        result = self.handle_call_tool(params)
                    else:
                        self.send_error_response(request_id, -32601, f"Method not found: {method}")
                        continue
                    
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": result
                    }
                    self.send_response(response)
                    
                except json.JSONDecodeError:
                    self.send_error_response(None, -32700, "Parse error")
                except Exception as e:
                    logger.error(f"Error in main loop: {e}")
                    self.send_error_response(request.get("id"), -32603, f"Internal error: {e}")
                    
        finally:
            db_handler.close_all_connections()
            logger.info("All database connections closed.")

if __name__ == "__main__":
    server = MCPServer()
    server.run()