import { CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { DatabaseManager } from '../database/manager.js';
import { QueryValidator } from '../utils/validator.js';
import { logger } from '../utils/logger.js';
import { z } from 'zod';
const QuerySchema = z.object({
sql: z.string().min(1, 'SQL query is required'),
params: z.array(z.any()).optional().default([]),
});
export async function handleQuery(request: z.infer<typeof CallToolRequestSchema>) {
try {
// Validate request
if (request.params.name !== 'mysql_query') {
throw new Error(`Unknown tool: ${request.params.name}`);
}
// Validate arguments
const { sql, params } = QuerySchema.parse(request.params.arguments);
// Validate query permissions
const validation = QueryValidator.validateQuery(sql);
if (!validation.allowed) {
return {
content: [
{
type: 'text',
text: `Query not allowed: ${validation.reason}`
}
]
};
}
// Execute query
const db = DatabaseManager.getInstance();
const result = await db.executeQuery(sql, params);
// Format response
let responseText = '';
if (result.rows.length > 0) {
responseText = `Query executed successfully. Retrieved ${result.rows.length} rows:\n\n`;
responseText += JSON.stringify(result.rows, null, 2);
} else if (result.affectedRows !== undefined) {
responseText = `Query executed successfully. ${result.affectedRows} rows affected.`;
if (result.insertId) {
responseText += ` Insert ID: ${result.insertId}`;
}
} else {
responseText = 'Query executed successfully. No rows returned.';
}
return {
content: [
{
type: 'text',
text: responseText
}
]
};
} catch (error) {
logger.error('Query execution failed', { error, request });
return {
content: [
{
type: 'text',
text: `Error executing query: ${error instanceof Error ? error.message : 'Unknown error'}`
}
]
};
}
}
