MCP MySQL Server

import logging from typing import Dict, Any, List, Optional, Tuple from .connection import get_db_manager logger = logging.getLogger(__name__) class QueryExecutor: def __init__(self): self.db = get_db_manager() def get_tables(self) -> List[Dict[str, Any]]: """Get all tables in the database""" query = """ SELECT TABLE_NAME as table_name, TABLE_TYPE as table_type, ENGINE as engine, TABLE_ROWS as row_count, DATA_LENGTH as data_length, INDEX_LENGTH as index_length, TABLE_COMMENT as comment FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() ORDER BY TABLE_NAME """ return self.db.execute_query(query) def get_table_columns(self, table_name: str) -> List[Dict[str, Any]]: """Get columns information for a specific table""" query = """ SELECT COLUMN_NAME as column_name, DATA_TYPE as data_type, IS_NULLABLE as is_nullable, COLUMN_DEFAULT as default_value, COLUMN_KEY as key_type, EXTRA as extra, COLUMN_COMMENT as comment FROM information_schema.COLUMNS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s ORDER BY ORDINAL_POSITION """ return self.db.execute_query(query, (table_name,)) def get_table_indexes(self, table_name: str) -> List[Dict[str, Any]]: """Get indexes for a specific table""" query = """ SELECT INDEX_NAME as index_name, COLUMN_NAME as column_name, NON_UNIQUE as non_unique, SEQ_IN_INDEX as sequence, INDEX_TYPE as index_type FROM information_schema.STATISTICS WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s ORDER BY INDEX_NAME, SEQ_IN_INDEX """ return self.db.execute_query(query, (table_name,)) def get_sample_data(self, table_name: str, limit: int = 10) -> List[Dict[str, Any]]: """Get sample data from a table""" # Sanitize table name to prevent SQL injection if not table_name.replace('_', '').replace('-', '').isalnum(): raise ValueError("Invalid table name") query = f"SELECT * FROM `{table_name}` LIMIT %s" return self.db.execute_query(query, (limit,)) def get_table_stats(self, table_name: str) -> Dict[str, Any]: """Get basic statistics for a table""" if not table_name.replace('_', '').replace('-', '').isalnum(): raise ValueError("Invalid table name") # Get row count count_query = f"SELECT COUNT(*) as row_count FROM `{table_name}`" row_count = self.db.execute_query(count_query)[0]['row_count'] # Get table info info_query = """ SELECT TABLE_NAME as table_name, ENGINE as engine, DATA_LENGTH as data_length, INDEX_LENGTH as index_length, AUTO_INCREMENT as auto_increment FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s """ table_info = self.db.execute_query(info_query, (table_name,)) if table_info: table_info[0]['row_count'] = row_count return table_info[0] return {'table_name': table_name, 'row_count': row_count} def search_tables_by_name(self, search_term: str) -> List[Dict[str, Any]]: """Search tables by name pattern""" query = """ SELECT TABLE_NAME as table_name, TABLE_TYPE as table_type, TABLE_COMMENT as comment FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME LIKE %s ORDER BY TABLE_NAME """ pattern = f"%{search_term}%" return self.db.execute_query(query, (pattern,)) def get_foreign_keys(self, table_name: str) -> List[Dict[str, Any]]: """Get foreign key relationships for a table""" query = """ SELECT COLUMN_NAME as column_name, REFERENCED_TABLE_NAME as referenced_table, REFERENCED_COLUMN_NAME as referenced_column, CONSTRAINT_NAME as constraint_name FROM information_schema.KEY_COLUMN_USAGE WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s AND REFERENCED_TABLE_NAME IS NOT NULL """ return self.db.execute_query(query, (table_name,)) def execute_custom_query(self, query: str, params: Optional[tuple] = None) -> List[Dict[str, Any]]: """Execute a custom SELECT query""" # Basic validation to ensure it's a SELECT query query_stripped = query.strip().upper() if not query_stripped.startswith('SELECT'): raise ValueError("Only SELECT queries are allowed") return self.db.execute_query(query, params) # Global query executor instance query_executor = QueryExecutor()