@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
JavaScript
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}]`;
}