MySQL Navigator MCP

# --- Standard Library Imports --- import sys import logging import json import re from pathlib import Path from typing import Dict, Optional import argparse # --- Third-party Library Imports --- from mysql.connector import Error from mcp.server.fastmcp import FastMCP # --- Local Application/Library Specific Imports --- from mcp_db.db import DatabaseCredentials from mcp_db.connection_pool import ConnectionPool from mcp_db.security import SecurityManager, enforce_security from mcp_db.metrics import MetricsTracker, track_query_metrics # ============================================================================== # Logging Configuration # ============================================================================== log_dir = Path.home() / ".mcp" log_dir.mkdir(parents=True, exist_ok=True) log_file = log_dir / "mcp-db.log" logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(sys.stdout), # Log to console logging.FileHandler(log_file, encoding='utf-8') ] ) logger = logging.getLogger(__name__) # ============================================================================== # Custom Error Class # ============================================================================== class MCPError(Exception): """Custom exception class for MCP-specific errors.""" def __init__(self, message: str, details: Optional[Dict] = None): self.message = message self.details = details or {} super().__init__(self.message) def to_dict(self) -> Dict: return {"error": self.message, "details": self.details} # ============================================================================== # Error Handling Decorator # ============================================================================== def handle_error(func): def wrapper(*args, **kwargs): try: return func(*args, **kwargs) except MCPError as e: logger.error(f"MCP Error in {func.__name__}: {e.message}", exc_info=False) logger.debug(f"MCP Error Details: {e.details}") try: return json.dumps(e.to_dict()) except Exception as serialization_error: logger.critical(f"Failed to serialize MCPError: {serialization_error}", exc_info=True) return json.dumps({"error": "Internal Server Error: Failed to report specific error."}) except Error as db_error: logger.error(f"Database Error in {func.__name__}: {db_error}", exc_info=True) error_dict = { "error": f"Database operation failed: {db_error.msg}", "details": { "type": type(db_error).__name__, "errno": db_error.errno, "sqlstate": db_error.sqlstate, } } try: return json.dumps(error_dict) except Exception as serialization_error: logger.critical(f"Failed to serialize Database Error details: {serialization_error}", exc_info=True) return json.dumps({"error": "Internal Server Error: Failed to report database error details."}) except Exception as e: logger.error(f"Unexpected Error in {func.__name__}: {e}", exc_info=True) error_dict = { "error": "An unexpected internal server error occurred.", "details": { "type": type(e).__name__, "message": str(e), } } try: return json.dumps(error_dict) except Exception as serialization_error: logger.critical(f"CRITICAL: Failed to serialize unexpected error details: {serialization_error}", exc_info=True) return json.dumps({"error": "Internal Server Error: Failed to report error details."}) wrapper.__name__ = func.__name__ wrapper.__doc__ = func.__doc__ wrapper.__annotations__ = func.__annotations__ return wrapper # ============================================================================== # Environment Loading # ============================================================================== def load_env_file(env_path: str) -> Dict[str, str]: if not env_path: logger.critical("No .env file path provided. Use --env to specify the path.") raise MCPError("Configuration Error: .env file path is required.") env_file = Path(env_path) if not env_file.is_file(): logger.critical(f".env file not found at specified path: {env_path}") raise MCPError(f"Configuration Error: .env file not found at {env_path}") env_vars = {} with open(env_file, "r") as f: for line in f: line = line.strip() if not line or line.startswith('#'): continue if '=' not in line: continue key, value = line.split('=', 1) env_vars[key.strip()] = value.strip() logger.info(f"Successfully loaded environment variables from {env_path}") return env_vars # ============================================================================== # Main Application Setup # ============================================================================== def run_server(env_file_path: str): try: env_vars = load_env_file(env_file_path) logger.info("Initializing MCP server...") mcp = FastMCP( "MySQL Navigator MCP", description="A powerful MySQL/MariaDB database navigation tool using MCP.", version="1.0.0" ) logger.info("MCP server framework initialized.") pool = ConnectionPool() security_manager = SecurityManager() metrics = MetricsTracker() logger.info("Core services (Connection Pool, Security, Metrics) initialized.") try: db_port = env_vars.get("DB_PORT") db_credentials = DatabaseCredentials( host=env_vars.get("DB_HOST", "localhost"), port=int(db_port) if db_port and db_port.isdigit() else 3306, database=env_vars.get("DB_NAME"), username=env_vars.get("DB_USER"), password=env_vars.get("DB_PASSWORD"), ssl_ca=env_vars.get("DB_SSL_CA"), max_retries=int(env_vars.get("DB_MAX_RETRIES", "3")), retry_delay=float(env_vars.get("DB_RETRY_DELAY", "1.0")) ) security_manager.sanitize_connection_params(db_credentials) logger.info(f"Database credentials loaded and validated for initial database: {db_credentials.database}") except Exception as e: logger.critical(f"Failed to load or validate initial database credentials from environment: {e}", exc_info=True) raise MCPError(f"Configuration Error: Invalid database credentials in .env: {e}") try: pool.initialize_pool(db_credentials) logger.info(f"Connection pool initialized for database: {db_credentials.database}") except Exception as e: logger.critical(f"Failed to initialize connection pool: {e}", exc_info=True) raise MCPError(f"Initialization Error: Could not establish initial database connection pool: {e}") # MCP Tool Definitions @mcp.tool() @handle_error def connect_to_database() -> str: global db_credentials logger.info(f"Attempting to connect/verify connection to database: {db_credentials.database}") try: pool.close_all() pool.initialize_pool(db_credentials) conn = pool.get_connection() conn.ping(reconnect=True, attempts=1, delay=0) conn.close() logger.info(f"Successfully connected/verified connection to database: {db_credentials.database}") status = {"status": "connected", "database": db_credentials.database, "host": db_credentials.host} return json.dumps(status) except Error as db_err: logger.error(f"Database connection/verification failed: {db_err}", exc_info=True) raise MCPError(f"Failed to connect to database '{db_credentials.database}'", {"details": str(db_err), "errno": db_err.errno}) except Exception as e: logger.error(f"Unexpected error during connect_to_database: {e}", exc_info=True) raise MCPError("Failed to connect to database due to an unexpected error.", {"details": str(e)}) @mcp.tool() @handle_error def switch_database(database: str) -> str: global db_credentials logger.info(f"Attempting to switch database connection to: {database}") if not database or not isinstance(database, str): raise MCPError("Invalid input: 'database' parameter must be a non-empty string.") try: new_credentials = DatabaseCredentials( host=db_credentials.host, port=db_credentials.port, database=database.strip(), username=db_credentials.username, password=db_credentials.password, ssl_ca=db_credentials.ssl_ca, max_retries=db_credentials.max_retries, retry_delay=db_credentials.retry_delay ) security_manager.sanitize_connection_params(new_credentials) pool.close_all() pool.initialize_pool(new_credentials) conn = pool.get_connection() conn.ping(reconnect=True, attempts=1, delay=0) conn.close() db_credentials = new_credentials logger.info(f"Successfully switched connection pool to database: {database}") status = {"status": "switched", "database": database, "host": db_credentials.host} return json.dumps(status) except Error as db_err: logger.error(f"Failed to switch to database '{database}': {db_err}", exc_info=True) raise MCPError(f"Failed to switch to database '{database}'", {"details": str(db_err), "errno": db_err.errno}) except Exception as e: logger.error(f"Unexpected error during switch_database: {e}", exc_info=True) raise MCPError(f"Failed to switch database due to an unexpected error.", {"details": str(e)}) @mcp.tool() @handle_error @enforce_security def load_database_schema() -> str: logger.info(f"Loading schema for database: {db_credentials.database}") connection = None try: connection = pool.get_connection() cursor = connection.cursor() cursor.execute("SHOW TABLES") tables = [table[0] for table in cursor.fetchall()] logger.debug(f"Found tables: {tables}") schema = {} for table_name in tables: if not re.match(r'^[a-zA-Z0-9_]+$', table_name): logger.warning(f"Skipping potentially unsafe table name for DESCRIBE: {table_name}") continue cursor.execute(f"DESCRIBE `{table_name}`") columns = cursor.fetchall() schema[table_name] = [{"name": col[0], "type": col[1]} for col in columns] logger.debug(f"Schema for table '{table_name}': {schema[table_name]}") logger.info(f"Successfully loaded schema for {len(schema)} tables.") return json.dumps({"schema": schema}) except Error as db_err: logger.error(f"Error getting database schema: {db_err}", exc_info=True) raise MCPError("Error getting database schema", {"details": str(db_err), "errno": db_err.errno}) except Exception as e: logger.error(f"Unexpected error loading schema: {e}", exc_info=True) raise MCPError("Unexpected error loading schema.", {"details": str(e)}) finally: if connection and connection.is_connected(): cursor.close() connection.close() logger.debug("Database connection closed after loading schema.") @mcp.tool() @handle_error @enforce_security @track_query_metrics def query_database(query_params: dict) -> str: logger.info(f"Received query request for database: {db_credentials.database}") logger.debug(f"Query parameters received: {query_params}") connection = None try: table_name = query_params.get("table_name") if not table_name or not isinstance(table_name, str) or not re.match(r'^[a-zA-Z0-9_]+$', table_name): raise MCPError("Invalid or missing 'table_name'. Must be alphanumeric/underscore.") select_fields_input = query_params.get("select_fields", ["*"]) if not isinstance(select_fields_input, list) or not select_fields_input: raise MCPError("'select_fields' must be a non-empty list.") select_fields = [] for field in select_fields_input: if field == "*": select_fields.append("*") elif isinstance(field, str) and re.match(r'^[a-zA-Z0-9_]+$', field): select_fields.append(f"`{field}`") else: raise MCPError(f"Invalid character in select field: {field}") select_clause = ", ".join(select_fields) where_conditions = query_params.get("where_conditions") order_by_input = query_params.get("order_by") order_direction_input = query_params.get("order_direction", "ASC").upper() limit_input = query_params.get("limit") offset_input = query_params.get("offset") query_parts = [f"SELECT {select_clause} FROM `{table_name}`"] parameters = [] if isinstance(where_conditions, dict) and where_conditions: where_clauses = [] for key, value in where_conditions.items(): if not isinstance(key, str) or not re.match(r'^[a-zA-Z0-9_]+$', key): raise MCPError(f"Invalid character in where condition key: {key}") where_clauses.append(f"`{key}` = %s") parameters.append(value) if where_clauses: query_parts.append("WHERE " + " AND ".join(where_clauses)) if isinstance(order_by_input, list) and order_by_input: order_by_clauses = [] for field in order_by_input: if not isinstance(field, str) or not re.match(r'^[a-zA-Z0-9_]+$', field): raise MCPError(f"Invalid character in order_by field: {field}") order_by_clauses.append(f"`{field}`") if order_by_clauses: if order_direction_input not in ["ASC", "DESC"]: raise MCPError("Invalid 'order_direction'. Must be 'ASC' or 'DESC'.") query_parts.append(f"ORDER BY {', '.join(order_by_clauses)} {order_direction_input}") effective_limit = 100 if limit_input is not None: try: limit_val = int(limit_input) if limit_val < 1: raise MCPError("'limit' must be a positive integer.") effective_limit = min(limit_val, 500) except (ValueError, TypeError): raise MCPError("'limit' must be an integer.") query_parts.append("LIMIT %s") parameters.append(effective_limit) if offset_input is not None: try: offset_val = int(offset_input) if offset_val < 0: raise MCPError("'offset' must be a non-negative integer.") query_parts.append("OFFSET %s") parameters.append(offset_val) except (ValueError, TypeError): raise MCPError("'offset' must be an integer.") final_query = " ".join(query_parts) logger.info(f"Executing parameterized query: {final_query}") logger.debug(f"Parameters: {tuple(parameters)}") connection = pool.get_connection() cursor = connection.cursor(dictionary=True) cursor.execute(final_query, tuple(parameters)) results = cursor.fetchall() logger.info(f"Query executed successfully. Retrieved {len(results)} rows.") return json.dumps({"results": results, "row_count": len(results)}) except Error as db_err: logger.error(f"Database query execution failed: {db_err}", exc_info=True) raise MCPError(f"Database query failed: {db_err.msg}", {"details": str(db_err), "errno": db_err.errno, "sqlstate": db_err.sqlstate}) except MCPError: raise except Exception as e: logger.error(f"Unexpected error during query execution: {e}", exc_info=True) raise MCPError("An unexpected error occurred while executing the query.", {"details": str(e)}) finally: if connection and connection.is_connected(): if 'cursor' in locals() and cursor: cursor.close() connection.close() logger.debug("Database connection closed after query.") logger.info("Starting MySQL Navigator MCP...") try: if pool._pool is None: logger.critical("Connection pool not available. Aborting server start.") sys.exit("Error: Database connection pool failed to initialize.") logger.info("MCP server setup complete. Ready to accept connections.") mcp.run() except KeyboardInterrupt: logger.info("Received KeyboardInterrupt (Ctrl+C). Initiating shutdown...") except MCPError as e: logger.critical(f"MCP Initialization Error: {e.message} - Details: {e.details}") sys.exit(f"Error: {e.message}") except Exception as e: logger.critical(f"Failed to start MCP server: {e}", exc_info=True) sys.exit(f"Error: Failed to start server - {e}") finally: logger.info("Initiating server shutdown sequence...") if 'pool' in locals() and pool: pool.close_all() logger.info("Database connection pool closed.") else: logger.warning("Connection pool object not found during shutdown.") logger.info("MySQL Navigator MCP server shutdown complete.") except Exception as e: logger.critical(f"A critical error occurred outside the main execution block: {e}", exc_info=True) sys.exit(f"Critical Error: {e}") def main_cli(): parser = argparse.ArgumentParser(description="MySQL Navigator MCP Server") parser.add_argument('--config', type=str, required=True, help="Path to .env file with DB credentials") args = parser.parse_args() run_server(args.config)