UNPKG

@bluspace/mssql-mcp-server

Version:

Microsoft SQL Server MCP (Model Context Protocol) Server - AI-powered database interaction tool

235 lines (234 loc) 10.6 kB
import { z } from "zod"; import sql from "mssql"; // Store for active transactions (in production, use a proper session manager) const activeTransactions = new Map(); export function registerTransactionTools(server, db) { server.tool("executeTransaction", "Execute multiple operations in a transaction", { operations: z.array(z.object({ tool: z.enum(["createRow", "updateRow", "deleteRow", "executeQuery"]), params: z.record(z.any()) })).describe("Operations to execute in order"), isolationLevel: z.enum([ "READ_UNCOMMITTED", "READ_COMMITTED", "REPEATABLE_READ", "SERIALIZABLE", "SNAPSHOT" ]).optional().default("READ_COMMITTED").describe("Transaction isolation level") }, async ({ operations, isolationLevel = "READ_COMMITTED" }) => { const transaction = new sql.Transaction(db); const results = []; try { // Map isolation level to SQL Server constants const isolationLevelMap = { "READ_UNCOMMITTED": sql.ISOLATION_LEVEL.READ_UNCOMMITTED, "READ_COMMITTED": sql.ISOLATION_LEVEL.READ_COMMITTED, "REPEATABLE_READ": sql.ISOLATION_LEVEL.REPEATABLE_READ, "SERIALIZABLE": sql.ISOLATION_LEVEL.SERIALIZABLE, "SNAPSHOT": sql.ISOLATION_LEVEL.SNAPSHOT }; // Begin transaction with specified isolation level await transaction.begin(isolationLevelMap[isolationLevel]); // Execute each operation for (let i = 0; i < operations.length; i++) { const op = operations[i]; const request = new sql.Request(transaction); try { let result; switch (op.tool) { case "createRow": { const { table, row } = op.params; const cols = Object.keys(row).map(k => `[${k}]`).join(","); const paramNames = Object.keys(row).map((_, idx) => `@param${idx}`).join(","); Object.entries(row).forEach(([_, value], idx) => { request.input(`param${idx}`, value); }); const insertResult = await request.query(`INSERT INTO ${formatTableName(table)}(${cols}) OUTPUT INSERTED.* VALUES(${paramNames})`); result = { inserted: insertResult.recordset[0] }; break; } case "updateRow": { const { table, key, changes } = op.params; const setClauses = []; let paramIdx = 0; Object.entries(changes).forEach(([k, v]) => { request.input(`set_${paramIdx}`, v); setClauses.push(`[${k}] = @set_${paramIdx}`); paramIdx++; }); const whereClauses = []; Object.entries(key).forEach(([k, v]) => { request.input(`where_${paramIdx}`, v); whereClauses.push(`[${k}] = @where_${paramIdx}`); paramIdx++; }); const updateResult = await request.query(`UPDATE ${formatTableName(table)} SET ${setClauses.join(",")} OUTPUT INSERTED.* WHERE ${whereClauses.join(" AND ")}`); result = { updated: updateResult.recordset[0] }; break; } case "deleteRow": { const { table, key } = op.params; const whereClauses = []; let paramIdx = 0; Object.entries(key).forEach(([k, v]) => { request.input(`param${paramIdx}`, v); whereClauses.push(`[${k}] = @param${paramIdx}`); paramIdx++; }); const deleteResult = await request.query(`DELETE FROM ${formatTableName(table)} OUTPUT DELETED.* WHERE ${whereClauses.join(" AND ")}`); result = { deleted: deleteResult.recordset[0] }; break; } case "executeQuery": { const { query, parameters = {} } = op.params; Object.entries(parameters).forEach(([name, value]) => { request.input(name, value); }); const queryResult = await request.query(query); result = { recordset: queryResult.recordset, rowsAffected: queryResult.rowsAffected }; break; } } results.push({ operation: i, tool: op.tool, success: true, result }); } catch (opError) { throw new Error(`Operation ${i} (${op.tool}) failed: ${opError.message}`); } } // Commit transaction await transaction.commit(); return { content: [{ type: "text", text: JSON.stringify({ success: true, operations: results, message: `Successfully executed ${operations.length} operations in transaction` }) }] }; } catch (error) { // Rollback on any error try { await transaction.rollback(); } catch (rollbackError) { console.error("Rollback failed:", rollbackError); } console.error("Transaction failed:", error); throw new Error(`Transaction failed: ${error.message}`); } }); // Advanced transaction control with session support server.tool("beginTransaction", "Begin a new transaction and return a transaction ID", { isolationLevel: z.enum([ "READ_UNCOMMITTED", "READ_COMMITTED", "REPEATABLE_READ", "SERIALIZABLE", "SNAPSHOT" ]).optional().default("READ_COMMITTED") }, async ({ isolationLevel = "READ_COMMITTED" }) => { try { const transaction = new sql.Transaction(db); const transactionId = `txn_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; const isolationLevelMap = { "READ_UNCOMMITTED": sql.ISOLATION_LEVEL.READ_UNCOMMITTED, "READ_COMMITTED": sql.ISOLATION_LEVEL.READ_COMMITTED, "REPEATABLE_READ": sql.ISOLATION_LEVEL.REPEATABLE_READ, "SERIALIZABLE": sql.ISOLATION_LEVEL.SERIALIZABLE, "SNAPSHOT": sql.ISOLATION_LEVEL.SNAPSHOT }; await transaction.begin(isolationLevelMap[isolationLevel]); activeTransactions.set(transactionId, transaction); return { content: [{ type: "text", text: JSON.stringify({ transactionId, status: "active", isolationLevel }) }] }; } catch (error) { console.error("Error beginning transaction:", error); throw new Error(`Failed to begin transaction: ${error.message}`); } }); server.tool("commitTransaction", "Commit an active transaction", { transactionId: z.string().describe("Transaction ID from beginTransaction") }, async ({ transactionId }) => { try { const transaction = activeTransactions.get(transactionId); if (!transaction) { throw new Error(`Transaction ${transactionId} not found or already completed`); } await transaction.commit(); activeTransactions.delete(transactionId); return { content: [{ type: "text", text: JSON.stringify({ transactionId, status: "committed" }) }] }; } catch (error) { console.error("Error committing transaction:", error); throw new Error(`Failed to commit transaction: ${error.message}`); } }); server.tool("rollbackTransaction", "Rollback an active transaction", { transactionId: z.string().describe("Transaction ID from beginTransaction") }, async ({ transactionId }) => { try { const transaction = activeTransactions.get(transactionId); if (!transaction) { throw new Error(`Transaction ${transactionId} not found or already completed`); } await transaction.rollback(); activeTransactions.delete(transactionId); return { content: [{ type: "text", text: JSON.stringify({ transactionId, status: "rolled back" }) }] }; } catch (error) { console.error("Error rolling back transaction:", error); throw new Error(`Failed to rollback transaction: ${error.message}`); } }); } // Helper function to format table name with schema function formatTableName(tableName) { const parts = tableName.split('.'); if (parts.length === 2) { return `[${parts[0]}].[${parts[1]}]`; } return `[dbo].[${tableName}]`; }