@bluspace/mssql-mcp-server
Version:
Microsoft SQL Server MCP (Model Context Protocol) Server - AI-powered database interaction tool
401 lines (399 loc) • 18.5 kB
JavaScript
import { z } from "zod";
import sql from "mssql";
export function registerBulkOperationTools(server, db) {
// Bulk insert
server.tool("bulkInsert", "Insert multiple rows efficiently", {
table: z.string().describe("Table name (can include schema)"),
rows: z.array(z.record(z.any())).describe("Array of row objects to insert"),
batchSize: z.number().optional().default(1000).describe("Number of rows per batch"),
validateSchema: z.boolean().optional().default(true).describe("Validate data against table schema"),
returnInserted: z.boolean().optional().default(false).describe("Return inserted rows (slower)")
}, async ({ table, rows, batchSize = 1000, validateSchema = true, returnInserted = false }) => {
if (rows.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({
inserted: 0,
batches: 0,
message: "No rows to insert"
})
}]
};
}
try {
// Get table schema if validation is requested
let columnInfo = [];
if (validateSchema) {
const schemaRequest = new sql.Request(db);
const tableParts = table.split('.');
const schemaName = tableParts.length === 2 ? tableParts[0] : 'dbo';
const tableName = tableParts.length === 2 ? tableParts[1] : table;
schemaRequest.input('schema', sql.NVarChar, schemaName);
schemaRequest.input('table', sql.NVarChar, tableName);
const schemaResult = await schemaRequest.query(`
SELECT
COLUMN_NAME as name,
DATA_TYPE as type,
IS_NULLABLE as nullable,
CHARACTER_MAXIMUM_LENGTH as maxLength
FROM INFORMATION_SCHEMA.COLUMNS
WHERE TABLE_SCHEMA = @schema AND TABLE_NAME = @table
`);
columnInfo = schemaResult.recordset;
}
// Process in batches
const totalBatches = Math.ceil(rows.length / batchSize);
let totalInserted = 0;
const insertedRows = [];
const errors = [];
for (let batchNum = 0; batchNum < totalBatches; batchNum++) {
const startIdx = batchNum * batchSize;
const endIdx = Math.min(startIdx + batchSize, rows.length);
const batchRows = rows.slice(startIdx, endIdx);
try {
// Use TVP (Table-Valued Parameters) for better performance if available
// For now, we'll use a transaction with multiple inserts
const transaction = new sql.Transaction(db);
await transaction.begin();
try {
for (const row of batchRows) {
const request = new sql.Request(transaction);
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);
});
let query = `INSERT INTO ${formatTableName(table)}(${cols}) VALUES(${paramNames})`;
if (returnInserted) {
query = `INSERT INTO ${formatTableName(table)}(${cols}) OUTPUT INSERTED.* VALUES(${paramNames})`;
}
const result = await request.query(query);
if (returnInserted && result.recordset.length > 0) {
insertedRows.push(result.recordset[0]);
}
totalInserted++;
}
await transaction.commit();
}
catch (batchError) {
await transaction.rollback();
errors.push({
batch: batchNum,
startRow: startIdx,
endRow: endIdx,
error: batchError.message
});
}
}
catch (error) {
errors.push({
batch: batchNum,
startRow: startIdx,
endRow: endIdx,
error: error.message
});
}
}
const response = {
inserted: totalInserted,
batches: totalBatches,
errors: errors
};
if (returnInserted) {
response.rows = insertedRows;
}
return {
content: [{
type: "text",
text: JSON.stringify(response)
}]
};
}
catch (error) {
console.error(`Error in bulk insert for ${table}:`, error);
throw new Error(`Bulk insert failed: ${error.message}`);
}
});
// Bulk update
server.tool("bulkUpdate", "Update multiple rows efficiently", {
table: z.string().describe("Table name (can include schema)"),
updates: z.array(z.object({
key: z.record(z.any()).describe("Primary key values to identify the row"),
changes: z.record(z.any()).describe("Column values to update")
})).describe("Array of update operations"),
batchSize: z.number().optional().default(1000).describe("Number of updates per batch")
}, async ({ table, updates, batchSize = 1000 }) => {
if (updates.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({
updated: 0,
batches: 0,
message: "No rows to update"
})
}]
};
}
try {
const totalBatches = Math.ceil(updates.length / batchSize);
let totalUpdated = 0;
const errors = [];
for (let batchNum = 0; batchNum < totalBatches; batchNum++) {
const startIdx = batchNum * batchSize;
const endIdx = Math.min(startIdx + batchSize, updates.length);
const batchUpdates = updates.slice(startIdx, endIdx);
const transaction = new sql.Transaction(db);
await transaction.begin();
try {
for (let i = 0; i < batchUpdates.length; i++) {
const update = batchUpdates[i];
const request = new sql.Request(transaction);
// Build SET clause
const setClauses = [];
let paramIdx = 0;
Object.entries(update.changes).forEach(([k, v]) => {
request.input(`set_${paramIdx}`, v);
setClauses.push(`[${k}] = @set_${paramIdx}`);
paramIdx++;
});
// Build WHERE clause
const whereClauses = [];
Object.entries(update.key).forEach(([k, v]) => {
request.input(`where_${paramIdx}`, v);
whereClauses.push(`[${k}] = @where_${paramIdx}`);
paramIdx++;
});
const result = await request.query(`UPDATE ${formatTableName(table)}
SET ${setClauses.join(",")}
WHERE ${whereClauses.join(" AND ")}`);
if (result.rowsAffected[0] > 0) {
totalUpdated++;
}
}
await transaction.commit();
}
catch (batchError) {
await transaction.rollback();
errors.push({
batch: batchNum,
startRow: startIdx,
endRow: endIdx,
error: batchError.message
});
}
}
return {
content: [{
type: "text",
text: JSON.stringify({
updated: totalUpdated,
batches: totalBatches,
errors: errors
})
}]
};
}
catch (error) {
console.error(`Error in bulk update for ${table}:`, error);
throw new Error(`Bulk update failed: ${error.message}`);
}
});
// Bulk delete
server.tool("bulkDelete", "Delete multiple rows efficiently", {
table: z.string().describe("Table name (can include schema)"),
keys: z.array(z.record(z.any())).describe("Array of primary key objects to delete"),
batchSize: z.number().optional().default(1000).describe("Number of deletes per batch")
}, async ({ table, keys, batchSize = 1000 }) => {
if (keys.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({
deleted: 0,
batches: 0,
message: "No rows to delete"
})
}]
};
}
try {
const totalBatches = Math.ceil(keys.length / batchSize);
let totalDeleted = 0;
const errors = [];
for (let batchNum = 0; batchNum < totalBatches; batchNum++) {
const startIdx = batchNum * batchSize;
const endIdx = Math.min(startIdx + batchSize, keys.length);
const batchKeys = keys.slice(startIdx, endIdx);
const transaction = new sql.Transaction(db);
await transaction.begin();
try {
for (const key of batchKeys) {
const request = new sql.Request(transaction);
// Build WHERE clause
const whereClauses = [];
let paramIdx = 0;
Object.entries(key).forEach(([k, v]) => {
request.input(`param${paramIdx}`, v);
whereClauses.push(`[${k}] = @param${paramIdx}`);
paramIdx++;
});
const result = await request.query(`DELETE FROM ${formatTableName(table)}
WHERE ${whereClauses.join(" AND ")}`);
if (result.rowsAffected[0] > 0) {
totalDeleted++;
}
}
await transaction.commit();
}
catch (batchError) {
await transaction.rollback();
errors.push({
batch: batchNum,
startRow: startIdx,
endRow: endIdx,
error: batchError.message
});
}
}
return {
content: [{
type: "text",
text: JSON.stringify({
deleted: totalDeleted,
batches: totalBatches,
errors: errors
})
}]
};
}
catch (error) {
console.error(`Error in bulk delete for ${table}:`, error);
throw new Error(`Bulk delete failed: ${error.message}`);
}
});
// Bulk upsert (merge)
server.tool("bulkUpsert", "Insert or update multiple rows using MERGE", {
table: z.string().describe("Table name (can include schema)"),
rows: z.array(z.record(z.any())).describe("Array of row objects to upsert"),
keyColumns: z.array(z.string()).describe("Columns that uniquely identify a row"),
updateColumns: z.array(z.string()).optional().describe("Columns to update on match (defaults to all non-key columns)"),
batchSize: z.number().optional().default(1000).describe("Number of rows per batch")
}, async ({ table, rows, keyColumns, updateColumns, batchSize = 1000 }) => {
if (rows.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({
upserted: 0,
inserted: 0,
updated: 0,
message: "No rows to upsert"
})
}]
};
}
try {
let totalInserted = 0;
let totalUpdated = 0;
const errors = [];
// If updateColumns not specified, use all columns except key columns
if (!updateColumns) {
const allColumns = Object.keys(rows[0]);
updateColumns = allColumns.filter(col => !keyColumns.includes(col));
}
const totalBatches = Math.ceil(rows.length / batchSize);
for (let batchNum = 0; batchNum < totalBatches; batchNum++) {
const startIdx = batchNum * batchSize;
const endIdx = Math.min(startIdx + batchSize, rows.length);
const batchRows = rows.slice(startIdx, endIdx);
try {
// Create a temporary table for the batch
const tempTableName = `#TempUpsert_${Date.now()}_${batchNum}`;
const transaction = new sql.Transaction(db);
await transaction.begin();
try {
const request = new sql.Request(transaction);
// Create temp table with same structure
const columns = Object.keys(batchRows[0]);
const createTempTableQuery = `
CREATE TABLE ${tempTableName} (
${columns.map(col => `[${col}] NVARCHAR(MAX)`).join(', ')}
)
`;
await request.query(createTempTableQuery);
// Insert data into temp table
for (const row of batchRows) {
const insertRequest = new sql.Request(transaction);
const cols = columns.map(c => `[${c}]`).join(',');
const paramNames = columns.map((_, idx) => `@param${idx}`).join(',');
columns.forEach((col, idx) => {
insertRequest.input(`param${idx}`, row[col]);
});
await insertRequest.query(`INSERT INTO ${tempTableName} (${cols}) VALUES (${paramNames})`);
}
// Perform MERGE operation
const mergeRequest = new sql.Request(transaction);
const joinConditions = keyColumns.map(col => `T.[${col}] = S.[${col}]`).join(' AND ');
const updateSet = updateColumns.map(col => `T.[${col}] = S.[${col}]`).join(', ');
const insertColumns = columns.map(c => `[${c}]`).join(', ');
const insertValues = columns.map(c => `S.[${c}]`).join(', ');
const mergeResult = await mergeRequest.query(`
MERGE ${formatTableName(table)} AS T
USING ${tempTableName} AS S
ON ${joinConditions}
WHEN MATCHED THEN
UPDATE SET ${updateSet}
WHEN NOT MATCHED THEN
INSERT (${insertColumns})
VALUES (${insertValues});
SELECT @@ROWCOUNT as affected;
`);
const affected = mergeResult.recordset[0].affected;
// Note: This is approximate - MERGE doesn't directly tell us inserts vs updates
totalInserted += Math.floor(affected / 2);
totalUpdated += Math.ceil(affected / 2);
await transaction.commit();
}
catch (mergeError) {
await transaction.rollback();
throw mergeError;
}
}
catch (batchError) {
errors.push({
batch: batchNum,
startRow: startIdx,
endRow: endIdx,
error: batchError.message
});
}
}
return {
content: [{
type: "text",
text: JSON.stringify({
upserted: totalInserted + totalUpdated,
inserted: totalInserted,
updated: totalUpdated,
batches: totalBatches,
errors: errors
})
}]
};
}
catch (error) {
console.error(`Error in bulk upsert for ${table}:`, error);
throw new Error(`Bulk upsert failed: ${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}]`;
}