feat: return affected row count for INSERT/UPDATE/DELETE

tianzhou · claude · tianzhou · commit a200027e5b8b · 2025-12-25T21:11:05.000-08:00
- Add rowCount to all database connectors (PostgreSQL, MySQL, MariaDB, SQLite, SQL Server) - Add extractAffectedRows utility for MySQL/MariaDB result parsing - Update execute-sql and custom-tool-handler to use rowCount for count field - Frontend now displays "N rows affected" for write operations instead of "No results returned" 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
diff --git a/frontend/src/api/tools.ts b/frontend/src/api/tools.ts
@@ -75,7 +75,8 @@ export async function executeTool(
 
   const rows = toolResult.data.rows;
   if (rows.length === 0) {
-    return { columns: [], rows: [], rowCount: 0 };
+    // For INSERT/UPDATE/DELETE, rows is empty but count reflects affected rows
+    return { columns: [], rows: [], rowCount: toolResult.data.count };
   }
 
   const columns = Object.keys(rows[0]);
diff --git a/frontend/src/components/tool/ResultsTable.tsx b/frontend/src/components/tool/ResultsTable.tsx
@@ -108,11 +108,15 @@ export function ResultsTable({ result, error, isLoading, executedSql, executionT
     );
   }
 
-  // No results
+  // No rows returned - could be empty SELECT or successful INSERT/UPDATE/DELETE
   if (result.rows.length === 0) {
     return (
       <div className="border border-border rounded-lg bg-card p-8 text-center">
-        <p className="text-muted-foreground text-sm">No results returned</p>
+        <p className="text-muted-foreground text-sm">
+          {result.rowCount > 0
+            ? `${result.rowCount} row${result.rowCount !== 1 ? 's' : ''} affected`
+            : 'No results returned'}
+        </p>
       </div>
     );
   }
diff --git a/src/connectors/interface.ts b/src/connectors/interface.ts
@@ -9,7 +9,7 @@ export type ConnectorType = "postgres" | "mysql" | "mariadb" | "sqlite" | "sqlse
  */
 export interface SQLResult {
   rows: any[];
-  [key: string]: any;
+  rowCount: number;
 }
 
 export interface TableColumn {
diff --git a/src/connectors/mariadb/index.ts b/src/connectors/mariadb/index.ts
@@ -14,7 +14,7 @@ import {
 import { SafeURL } from "../../utils/safe-url.js";
 import { obfuscateDSNPassword } from "../../utils/dsn-obfuscate.js";
 import { SQLRowLimiter } from "../../utils/sql-row-limiter.js";
-import { parseQueryResults } from "../../utils/multi-statement-result-parser.js";
+import { parseQueryResults, extractAffectedRows } from "../../utils/multi-statement-result-parser.js";
 
 /**
  * MariaDB DSN Parser
@@ -542,7 +542,8 @@ export class MariaDBConnector implements Connector {
 
       // Parse results using shared utility that handles both single and multi-statement queries
       const rows = parseQueryResults(results);
-      return { rows };
+      const rowCount = extractAffectedRows(results);
+      return { rows, rowCount };
     } catch (error) {
       console.error("Error executing query:", error);
       throw error;
diff --git a/src/connectors/mysql/index.ts b/src/connectors/mysql/index.ts
@@ -14,7 +14,7 @@ import {
 import { SafeURL } from "../../utils/safe-url.js";
 import { obfuscateDSNPassword } from "../../utils/dsn-obfuscate.js";
 import { SQLRowLimiter } from "../../utils/sql-row-limiter.js";
-import { parseQueryResults } from "../../utils/multi-statement-result-parser.js";
+import { parseQueryResults, extractAffectedRows } from "../../utils/multi-statement-result-parser.js";
 
 /**
  * MySQL DSN Parser
@@ -553,7 +553,8 @@ export class MySQLConnector implements Connector {
 
       // Parse results using shared utility that handles both single and multi-statement queries
       const rows = parseQueryResults(firstResult);
-      return { rows };
+      const rowCount = extractAffectedRows(firstResult);
+      return { rows, rowCount };
     } catch (error) {
       console.error("Error executing query:", error);
       throw error;
diff --git a/src/connectors/postgres/index.ts b/src/connectors/postgres/index.ts
@@ -444,17 +444,21 @@ export class PostgresConnector implements Connector {
         const processedStatement = SQLRowLimiter.applyMaxRows(statements[0], options.maxRows);
 
         // Use parameters if provided
+        let result;
         if (parameters && parameters.length > 0) {
           try {
-            return await client.query(processedStatement, parameters);
+            result = await client.query(processedStatement, parameters);
           } catch (error) {
             console.error(`[PostgreSQL executeSQL] ERROR: ${(error as Error).message}`);
             console.error(`[PostgreSQL executeSQL] SQL: ${processedStatement}`);
             console.error(`[PostgreSQL executeSQL] Parameters: ${JSON.stringify(parameters)}`);
             throw error;
           }
+        } else {
+          result = await client.query(processedStatement);
         }
-        return await client.query(processedStatement);
+        // Explicitly return rows and rowCount to ensure rowCount is preserved
+        return { rows: result.rows, rowCount: result.rowCount ?? result.rows.length };
       } else {
         // Multiple statements - parameters not supported for multi-statement queries
         if (parameters && parameters.length > 0) {
@@ -463,6 +467,7 @@ export class PostgresConnector implements Connector {
 
         // Execute all in same session for transaction consistency
         let allRows: any[] = [];
+        let totalRowCount = 0;
 
         // Execute within a transaction to ensure session consistency
         await client.query('BEGIN');
@@ -476,14 +481,18 @@ export class PostgresConnector implements Connector {
             if (result.rows && result.rows.length > 0) {
               allRows.push(...result.rows);
             }
+            // Accumulate rowCount for INSERT/UPDATE/DELETE statements
+            if (result.rowCount) {
+              totalRowCount += result.rowCount;
+            }
           }
           await client.query('COMMIT');
         } catch (error) {
           await client.query('ROLLBACK');
           throw error;
         }
 
-        return { rows: allRows };
+        return { rows: allRows, rowCount: totalRowCount };
       }
     } finally {
       client.release();
diff --git a/src/connectors/sqlite/index.ts b/src/connectors/sqlite/index.ts
@@ -413,7 +413,7 @@ export class SQLiteConnector implements Connector {
           if (parameters && parameters.length > 0) {
             try {
               const rows = this.db.prepare(processedStatement).all(...parameters);
-              return { rows };
+              return { rows, rowCount: rows.length };
             } catch (error) {
               console.error(`[SQLite executeSQL] ERROR: ${(error as Error).message}`);
               console.error(`[SQLite executeSQL] SQL: ${processedStatement}`);
@@ -422,23 +422,24 @@ export class SQLiteConnector implements Connector {
             }
           } else {
             const rows = this.db.prepare(processedStatement).all();
-            return { rows };
+            return { rows, rowCount: rows.length };
           }
         } else {
           // Use run() for statements that don't return data
+          let result;
           if (parameters && parameters.length > 0) {
             try {
-              this.db.prepare(processedStatement).run(...parameters);
+              result = this.db.prepare(processedStatement).run(...parameters);
             } catch (error) {
               console.error(`[SQLite executeSQL] ERROR: ${(error as Error).message}`);
               console.error(`[SQLite executeSQL] SQL: ${processedStatement}`);
               console.error(`[SQLite executeSQL] Parameters: ${JSON.stringify(parameters)}`);
               throw error;
             }
           } else {
-            this.db.prepare(processedStatement).run();
+            result = this.db.prepare(processedStatement).run();
           }
-          return { rows: [] };
+          return { rows: [], rowCount: result.changes };
         }
       } else {
         // Multiple statements - parameters not supported for multi-statement queries
@@ -469,9 +470,11 @@ export class SQLiteConnector implements Connector {
           }
         }
 
-        // Execute write statements using native .exec() for optimal performance
-        if (writeStatements.length > 0) {
-          this.db.exec(writeStatements.join('; '));
+        // Execute write statements individually to track changes
+        let totalChanges = 0;
+        for (const statement of writeStatements) {
+          const result = this.db.prepare(statement).run();
+          totalChanges += result.changes;
         }
 
         // Execute read statements individually to collect results
@@ -483,7 +486,8 @@ export class SQLiteConnector implements Connector {
           allRows.push(...result);
         }
 
-        return { rows: allRows };
+        // rowCount is total changes for writes, plus rows returned for reads
+        return { rows: allRows, rowCount: totalChanges + allRows.length };
       }
     } catch (error) {
       throw error;
diff --git a/src/connectors/sqlserver/index.ts b/src/connectors/sqlserver/index.ts
@@ -556,12 +556,6 @@ export class SQLServerConnector implements Connector {
 
       return {
         rows: result.recordset || [],
-        fields:
-          result.recordset && result.recordset.length > 0
-            ? Object.keys(result.recordset[0]).map((key) => ({
-                name: key,
-              }))
-            : [],
         rowCount: result.rowsAffected[0] || 0,
       };
     } catch (error) {
diff --git a/src/tools/__tests__/execute-sql.test.ts b/src/tools/__tests__/execute-sql.test.ts
@@ -53,7 +53,7 @@ describe('execute-sql tool', () => {
 
   describe('basic execution', () => {
     it('should execute SELECT and return rows', async () => {
-      const mockResult: SQLResult = { rows: [{ id: 1, name: 'test' }] };
+      const mockResult: SQLResult = { rows: [{ id: 1, name: 'test' }], rowCount: 1 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const handler = createExecuteSqlToolHandler('test_source');
@@ -67,7 +67,7 @@ describe('execute-sql tool', () => {
     });
 
     it('should pass multi-statement SQL directly to connector', async () => {
-      const mockResult: SQLResult = { rows: [{ id: 1 }] };
+      const mockResult: SQLResult = { rows: [{ id: 1 }], rowCount: 1 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const sql = 'SELECT * FROM users; SELECT * FROM roles;';
@@ -102,7 +102,7 @@ describe('execute-sql tool', () => {
     });
 
     it('should allow SELECT statements', async () => {
-      const mockResult: SQLResult = { rows: [{ id: 1 }] };
+      const mockResult: SQLResult = { rows: [{ id: 1 }], rowCount: 1 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const handler = createExecuteSqlToolHandler('test_source');
@@ -114,7 +114,7 @@ describe('execute-sql tool', () => {
     });
 
     it('should allow multiple read-only statements', async () => {
-      const mockResult: SQLResult = { rows: [] };
+      const mockResult: SQLResult = { rows: [], rowCount: 0 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const sql = 'SELECT * FROM users; SELECT * FROM roles;';
@@ -173,7 +173,7 @@ describe('execute-sql tool', () => {
       mockGetToolRegistry.mockReturnValue({
         getBuiltinToolConfig: vi.fn().mockReturnValue(toolConfig),
       } as any);
-      const mockResult: SQLResult = { rows: [] };
+      const mockResult: SQLResult = { rows: [], rowCount: 0 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const handler = createExecuteSqlToolHandler('writable_source');
@@ -208,7 +208,7 @@ describe('execute-sql tool', () => {
       ['inline comments', 'SELECT id, -- user id\n       name FROM users'],
       ['only comments', '-- Just a comment\n/* Another */'],
     ])('should allow SELECT with %s', async (_, sql) => {
-      const mockResult: SQLResult = { rows: [] };
+      const mockResult: SQLResult = { rows: [], rowCount: 0 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const handler = createExecuteSqlToolHandler('test_source');
@@ -231,7 +231,7 @@ describe('execute-sql tool', () => {
       ['empty string', ''],
       ['only semicolons and whitespace', '   ;  ;  ; '],
     ])('should handle %s', async (_, sql) => {
-      const mockResult: SQLResult = { rows: [] };
+      const mockResult: SQLResult = { rows: [], rowCount: 0 };
       vi.mocked(mockConnector.executeSQL).mockResolvedValue(mockResult);
 
       const handler = createExecuteSqlToolHandler('test_source');
diff --git a/src/tools/custom-tool-handler.ts b/src/tools/custom-tool-handler.ts
@@ -199,7 +199,7 @@ export function createCustomToolHandler(toolConfig: ToolConfig) {
       // 7. Build response data
       const responseData = {
         rows: result.rows,
-        count: result.rows.length,
+        count: result.rowCount,
         source_id: toolConfig.source,
       };
 
diff --git a/src/tools/execute-sql.ts b/src/tools/execute-sql.ts
@@ -78,7 +78,7 @@ export function createExecuteSqlToolHandler(sourceId?: string) {
       // Build response data
       const responseData = {
         rows: result.rows,
-        count: result.rows.length,
+        count: result.rowCount,
         source_id: effectiveSourceId,
       };
 
diff --git a/src/utils/multi-statement-result-parser.ts b/src/utils/multi-statement-result-parser.ts
@@ -68,6 +68,43 @@ export function extractRowsFromMultiStatement(results: any): any[] {
   return allRows;
 }
 
+/**
+ * Extracts total affected rows from query results.
+ *
+ * For INSERT/UPDATE/DELETE operations, returns the sum of affectedRows.
+ * For SELECT operations, returns the number of rows.
+ *
+ * @param results - Raw results from the database driver
+ * @returns Total number of affected/returned rows
+ */
+export function extractAffectedRows(results: any): number {
+  // Handle metadata object (single INSERT/UPDATE/DELETE)
+  if (isMetadataObject(results)) {
+    return results.affectedRows || 0;
+  }
+
+  // Handle non-array results
+  if (!Array.isArray(results)) {
+    return 0;
+  }
+
+  // Check if this is a multi-statement result
+  if (isMultiStatementResult(results)) {
+    let totalAffected = 0;
+    for (const result of results) {
+      if (isMetadataObject(result)) {
+        totalAffected += result.affectedRows || 0;
+      } else if (Array.isArray(result)) {
+        totalAffected += result.length;
+      }
+    }
+    return totalAffected;
+  }
+
+  // Single statement result - results is the rows array directly
+  return results.length;
+}
+
 /**
  * Parses database query results, handling both single and multi-statement queries.
  *

Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,8 @@ export async function executeTool(`
`75`	`75`
`76`	`76`	`const rows = toolResult.data.rows;`
`77`	`77`	`if (rows.length === 0) {`
`78`		`- return { columns: [], rows: [], rowCount: 0 };`
	`78`	`+ // For INSERT/UPDATE/DELETE, rows is empty but count reflects affected rows`
	`79`	`+ return { columns: [], rows: [], rowCount: toolResult.data.count };`
`79`	`80`	`}`
`80`	`81`
`81`	`82`	`const columns = Object.keys(rows[0]);`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ export type ConnectorType = "postgres" \| "mysql" \| "mariadb" \| "sqlite" \| "sqlse`
`9`	`9`	`*/`
`10`	`10`	`export interface SQLResult {`
`11`	`11`	`rows: any[];`
`12`		`- [key: string]: any;`
	`12`	`+ rowCount: number;`
`13`	`13`	`}`
`14`	`14`
`15`	`15`	`export interface TableColumn {`