"""
MCP tools for schema operations
"""
import logging
from typing import Dict, Any
from services.schema_service import SchemaService
logger = logging.getLogger(__name__)
class SchemaTools:
"""
MCP tool implementations for schema operations
"""
def __init__(self, schema_service: SchemaService):
self.schema_service = schema_service
async def get_schema_info(self, include_sample_data: bool = False) -> Dict[str, Any]:
"""
MCP tool: Get database schema information
Args:
include_sample_data: Whether to include sample data
Returns:
Schema information
"""
try:
logger.info("MCP Tool - Get Schema Info")
schema_info = self.schema_service.get_schema_info()
# Convert to serializable format
result = {
'success': True,
'summary': schema_info.summary,
'table_count': len(schema_info.tables),
'relationship_count': len(schema_info.relationships),
'tables': []
}
for table in schema_info.tables:
table_info = {
'name': table['table_name'],
'columns': [
{
'name': col['name'],
'type': str(col['type']),
'nullable': col['nullable']
}
for col in table['columns']
],
'row_count': table.get('row_count', 0),
'primary_keys': table.get('primary_keys', {}).get('constrained_columns', []),
'foreign_keys': len(table.get('foreign_keys', []))
}
if include_sample_data:
table_info['sample_data'] = self.schema_service._get_sample_data(table['table_name'])
result['tables'].append(table_info)
return result
except Exception as e:
logger.error(f"Get schema info tool failed: {e}")
return {
'success': False,
'error': str(e),
'table_count': 0,
'relationship_count': 0,
'tables': []
}
async def find_relevant_tables(self, question: str) -> Dict[str, Any]:
"""
MCP tool: Find tables relevant to a question
Args:
question: Natural language question
Returns:
Relevant table information
"""
try:
logger.info(f"MCP Tool - Find Relevant Tables: {question}")
relevant_tables = self.schema_service.find_relevant_tables(question)
return {
'success': True,
'question': question,
'relevant_tables': relevant_tables,
'count': len(relevant_tables)
}
except Exception as e:
logger.error(f"Find relevant tables tool failed: {e}")
return {
'success': False,
'error': str(e),
'relevant_tables': [],
'count': 0
}
async def get_table_details(self, table_name: str) -> Dict[str, Any]:
"""
MCP tool: Get detailed information about a specific table
Args:
table_name: Name of the table
Returns:
Detailed table information
"""
try:
logger.info(f"MCP Tool - Get Table Details: {table_name}")
table_details = self.schema_service.get_table_details(table_name)
return {
'success': True,
'table_name': table_name,
'details': table_details
}
except Exception as e:
logger.error(f"Get table details tool failed: {e}")
return {
'success': False,
'error': str(e),
'table_name': table_name,
'details': {}
}
