execute_sql.ts•2.77 kB
import { z } from "zod";
import * as pg from "pg";
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
export interface ToolProps {
pool: pg.Pool;
}
export function ExecuteSQL(server: McpServer, props: ToolProps) {
server.tool(
"execute_sql",
"Execute any raw SQL statement for data queries or updates",
{
sql: z.string().describe("The SQL statement to execute"),
read_only: z
.boolean()
.optional()
.describe("Whether to execute in read-only mode (default: false)"),
},
async ({ sql, read_only = false }) => {
if (!props.pool) {
throw new Error("Database pool not initialized");
}
const client = await props.pool.connect();
try {
if (read_only) {
await client.query("BEGIN TRANSACTION READ ONLY");
} else {
await client.query("BEGIN");
}
const result = await client.query(sql);
if (read_only) {
await client.query("ROLLBACK");
} else {
await client.query("COMMIT");
}
// Determine the type of operation
const upperSQL = sql.trim().toUpperCase();
const isSelect = upperSQL.startsWith("SELECT") || upperSQL.startsWith("WITH");
const isInsert = upperSQL.startsWith("INSERT");
const isUpdate = upperSQL.startsWith("UPDATE");
const isDelete = upperSQL.startsWith("DELETE");
const isDDL = ["CREATE", "ALTER", "DROP", "TRUNCATE"].some((keyword) =>
upperSQL.startsWith(keyword),
);
let operationType = "UNKNOWN";
if (isSelect) operationType = "SELECT";
else if (isInsert) operationType = "INSERT";
else if (isUpdate) operationType = "UPDATE";
else if (isDelete) operationType = "DELETE";
else if (isDDL) operationType = "DDL";
const response: any = {
success: true,
operation_type: operationType,
rows_affected: result.rowCount || 0,
read_only_mode: read_only,
};
// Include rows for SELECT queries or if there are returned rows
if (isSelect || (result.rows && result.rows.length > 0)) {
response.rows = result.rows;
response.row_count = result.rows.length;
}
// Include field information for SELECT queries
if (isSelect && result.fields) {
response.fields = result.fields.map((field) => ({
name: field.name,
dataTypeID: field.dataTypeID,
dataTypeSize: field.dataTypeSize,
dataTypeModifier: field.dataTypeModifier,
}));
}
return {
content: [
{
type: "text",
text: JSON.stringify(response, null, 2),
},
],
};
} catch (error) {
await client.query("ROLLBACK").catch(() => {});
throw new Error(
`SQL execution failed: ${error instanceof Error ? error.message : "Unknown error"}`,
);
} finally {
client.release();
}
},
);
}