Allow different columns for ReadData and UpdateData

Author: Akshay-Mata

MssqlMcp/NodejsMcpServer/src/tools/ReadDataTool.ts

@@ -4,33 +4,42 @@ import { Tool } from "@modelcontextprotocol/sdk/types.js";
 export class ReadDataTool implements Tool {
   [key: string]: any;
   name = "read_data";
-  description = "Retrieves data from an MSSQL Database table by its ID";
+  description = "Executes a SELECT query on an MSSQL Database table. The query must start with SELECT for security.";
   inputSchema = {
     type: "object",
     properties: {
-      tableName: { type: "string", description: "Name of the table" },
-      id: { type: "string", description: "ID of the data to retrieve" },
+      query: { 
+        type: "string", 
+        description: "SQL SELECT query to execute (must start with SELECT). Example: SELECT * FROM movies WHERE genre = 'comedy'" 
+      },
     },
-    required: ["tableName", "id"],
+    required: ["query"],
   } as any;
 
   async run(params: any) {
     try {
-      const { tableName, id } = params;
+      const { query } = params;
+      
+      // Basic validation: ensure query starts with SELECT (case insensitive)
+      const trimmedQuery = query.trim();
+      if (!trimmedQuery.toUpperCase().startsWith('SELECT')) {
+        throw new Error("Query must start with SELECT for security reasons");
+      }
+
       const request = new sql.Request();
-      request.input("id", sql.VarChar, id);
-      const query = `SELECT * FROM ${tableName} WHERE id = @id`;
-      const result = await request.query(query);
+      const result = await request.query(trimmedQuery);
+      
       return {
         success: true,
-        message: `Data retrieved successfully`,
-        data: result.recordset[0],
+        message: `Query executed successfully. Retrieved ${result.recordset.length} record(s)`,
+        data: result.recordset,
+        recordCount: result.recordset.length,
       };
     } catch (error) {
-      console.error("Error retrieving data:", error);
+      console.error("Error executing query:", error);
       return {
         success: false,
-        message: `Failed to retrieve data: ${error}`,
+        message: `Failed to execute query: ${error}`,
       };
     }
   }

MssqlMcp/NodejsMcpServer/src/tools/UpdateDataTool.ts

@@ -4,36 +4,53 @@ import { Tool } from "@modelcontextprotocol/sdk/types.js";
 export class UpdateDataTool implements Tool {
   [key: string]: any;
   name = "update_data";
-  description = "Updates data in an MSSQL Database table by its ID";
+  description = "Updates data in an MSSQL Database table using a WHERE clause. The WHERE clause must be provided for security.";
   inputSchema = {
     type: "object",
     properties: {
-      tableName: { type: "string", description: "Name of the table" },
-      id: { type: "string", description: "ID of the data to update" },
+      tableName: { 
+        type: "string", 
+        description: "Name of the table to update" 
+      },
       updates: {
         type: "object",
-        description: "Key-value pairs of columns to update",
+        description: "Key-value pairs of columns to update. Example: { 'status': 'active', 'last_updated': '2025-01-01' }",
+      },
+      whereClause: { 
+        type: "string", 
+        description: "WHERE clause to identify which records to update. Example: \"genre = 'comedy' AND created_date <= '2025-07-05'\"" 
       },
     },
-    required: ["tableName", "id", "updates"],
+    required: ["tableName", "updates", "whereClause"],
   } as any;
 
   async run(params: any) {
     try {
-      const { tableName, id, updates } = params;
+      const { tableName, updates, whereClause } = params;
+      
+      // Basic validation: ensure whereClause is not empty
+      if (!whereClause || whereClause.trim() === '') {
+        throw new Error("WHERE clause is required for security reasons");
+      }
+
       const request = new sql.Request();
-      request.input("id", sql.VarChar, id);
+      
+      // Build SET clause with parameterized queries for security
       const setClause = Object.keys(updates)
-        .map((key) => {
-          request.input(key, updates[key]);
-          return `${key} = @${key}`;
+        .map((key, index) => {
+          const paramName = `update_${index}`;
+          request.input(paramName, updates[key]);
+          return `[${key}] = @${paramName}`;
         })
         .join(", ");
-      const query = `UPDATE ${tableName} SET ${setClause} WHERE id = @id`;
-      await request.query(query);
+
+      const query = `UPDATE [${tableName}] SET ${setClause} WHERE ${whereClause}`;
+      const result = await request.query(query);
+      
       return {
         success: true,
-        message: `Data updated successfully`,
+        message: `Update completed successfully. ${result.rowsAffected[0]} row(s) affected`,
+        rowsAffected: result.rowsAffected[0],
       };
     } catch (error) {
       console.error("Error updating data:", error);