@tejasanik/postgres-mcp-server 2.1.0 → 2.2.0

package/dist/tools/sql-tools.js

@@ -4,169 +4,17 @@ import * as path from 'path';
 import * as os from 'os';
 import { v4 as uuidv4 } from 'uuid';
 import { validateIdentifier, validateIndexType, isReadOnlySql, validatePositiveInteger } from '../utils/validation.js';
-const MAX_OUTPUT_CHARS = 50000; // Maximum characters before writing to file
-const MAX_ROWS_DEFAULT = 1000; // Default max rows in direct response
-const MAX_ROWS_LIMIT = 100000; // Absolute maximum rows
-const DEFAULT_SQL_LENGTH_LIMIT = 100000; // Default SQL query length limit (100KB)
-const MAX_PARAMS = 100; // Maximum number of query parameters
-const MAX_SQL_FILE_SIZE = 50 * 1024 * 1024; // Maximum SQL file size (50MB)
-const MAX_DRY_RUN_SAMPLE_ROWS = 10; // Maximum sample rows to return in dry-run
-/**
- * Extract detailed error information from a PostgreSQL error.
- * Captures all available fields to help AI quickly identify and fix issues.
- */
-function extractDryRunError(error) {
-    const result = {
-        message: error instanceof Error ? error.message : String(error)
-    };
-    // PostgreSQL errors have additional properties
-    if (error && typeof error === 'object') {
-        const pgError = error;
-        if (pgError.code)
-            result.code = String(pgError.code);
-        if (pgError.severity)
-            result.severity = String(pgError.severity);
-        if (pgError.detail)
-            result.detail = String(pgError.detail);
-        if (pgError.hint)
-            result.hint = String(pgError.hint);
-        if (pgError.position)
-            result.position = Number(pgError.position);
-        if (pgError.internalPosition)
-            result.internalPosition = Number(pgError.internalPosition);
-        if (pgError.internalQuery)
-            result.internalQuery = String(pgError.internalQuery);
-        if (pgError.where)
-            result.where = String(pgError.where);
-        if (pgError.schema)
-            result.schema = String(pgError.schema);
-        if (pgError.table)
-            result.table = String(pgError.table);
-        if (pgError.column)
-            result.column = String(pgError.column);
-        if (pgError.dataType)
-            result.dataType = String(pgError.dataType);
-        if (pgError.constraint)
-            result.constraint = String(pgError.constraint);
-        if (pgError.file)
-            result.file = String(pgError.file);
-        if (pgError.line)
-            result.line = String(pgError.line);
-        if (pgError.routine)
-            result.routine = String(pgError.routine);
-    }
-    return result;
-}
-/**
- * Check if a SQL statement contains operations that cannot be fully rolled back
- * or have side effects even within a transaction.
- */
-function detectNonRollbackableOperations(sql, statementIndex, lineNumber) {
-    const warnings = [];
-    const upperSql = sql.toUpperCase().trim();
-    // Operations that cannot run inside a transaction at all - MUST SKIP
-    if (upperSql.match(/\bVACUUM\b/)) {
-        warnings.push({
-            operation: 'VACUUM',
-            message: 'VACUUM cannot run inside a transaction block. Statement skipped.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    if (upperSql.match(/\bCLUSTER\b/) && !upperSql.includes('CREATE')) {
-        warnings.push({
-            operation: 'CLUSTER',
-            message: 'CLUSTER cannot run inside a transaction block. Statement skipped.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    if (upperSql.match(/\bREINDEX\b.*\bCONCURRENTLY\b/)) {
-        warnings.push({
-            operation: 'REINDEX_CONCURRENTLY',
-            message: 'REINDEX CONCURRENTLY cannot run inside a transaction block. Statement skipped.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    if (upperSql.match(/\bCREATE\s+INDEX\b.*\bCONCURRENTLY\b/)) {
-        warnings.push({
-            operation: 'CREATE_INDEX_CONCURRENTLY',
-            message: 'CREATE INDEX CONCURRENTLY cannot run inside a transaction block. Statement skipped.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    if (upperSql.match(/\bCREATE\s+DATABASE\b/)) {
-        warnings.push({
-            operation: 'CREATE_DATABASE',
-            message: 'CREATE DATABASE cannot run inside a transaction block. Statement skipped.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    if (upperSql.match(/\bDROP\s+DATABASE\b/)) {
-        warnings.push({
-            operation: 'DROP_DATABASE',
-            message: 'DROP DATABASE cannot run inside a transaction block. Statement skipped.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    // Operations that have side effects even when rolled back - MUST SKIP
-    if (upperSql.match(/\bNEXTVAL\s*\(/)) {
-        warnings.push({
-            operation: 'SEQUENCE',
-            message: 'NEXTVAL increments sequence even when transaction is rolled back. Statement skipped to prevent sequence consumption.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    if (upperSql.match(/\bSETVAL\s*\(/)) {
-        warnings.push({
-            operation: 'SEQUENCE',
-            message: 'SETVAL modifies sequence. Statement skipped to prevent side effects.',
-            statementIndex,
-            lineNumber,
-            mustSkip: true
-        });
-    }
-    // INSERT with SERIAL/BIGSERIAL columns may consume sequence values - WARNING ONLY (not skipped)
-    if (upperSql.match(/\bINSERT\s+INTO\b/)) {
-        warnings.push({
-            operation: 'SEQUENCE',
-            message: 'INSERT may consume sequence values (for SERIAL/BIGSERIAL columns) even when rolled back.',
-            statementIndex,
-            lineNumber,
-            mustSkip: false // Warning only, do not skip
-        });
-    }
-    // NOTIFY only sends on commit, so safe in dry-run (rollback prevents notification)
-    if (upperSql.match(/\bNOTIFY\b/)) {
-        warnings.push({
-            operation: 'NOTIFY',
-            message: 'NOTIFY sends notifications on commit. Since dry-run rolls back, notifications will NOT be sent.',
-            statementIndex,
-            lineNumber,
-            mustSkip: false // Safe to execute in dry-run since we rollback
-        });
-    }
-    return warnings;
-}
+// Import utilities from modular structure
+import { MAX_OUTPUT_CHARS, MAX_ROWS_DEFAULT, MAX_ROWS_LIMIT, DEFAULT_SQL_LENGTH_LIMIT, MAX_PARAMS, MAX_SQL_FILE_SIZE, MAX_DRY_RUN_SAMPLE_ROWS, MAX_BATCH_QUERIES, MAX_MUTATION_SAMPLE_SIZE, DEFAULT_MUTATION_SAMPLE_SIZE, MAX_HYPOTHETICAL_INDEXES, MAX_DRY_RUN_STATEMENTS, DEFAULT_DRY_RUN_STATEMENTS, MAX_PREVIEW_STATEMENTS, DEFAULT_PREVIEW_STATEMENTS, MAX_ROWS_PER_STATEMENT, SQL_TRUNCATION_SHORT, SQL_TRUNCATION_LONG, MAX_TABLES_TO_ANALYZE, } from './sql/utils/constants.js';
+import { extractDryRunError, detectNonRollbackableOperations, } from './sql/utils/dry-run-utils.js';
+import { splitSqlStatementsWithLineNumbers, stripLeadingComments, detectStatementType, extractTablesFromSql, } from './sql/utils/sql-parser.js';
+import { preprocessSqlContent, } from './sql/utils/file-handler.js';
+// Connection utilities are handled inline for explicit control flow
+import { calculateExecutionTime, getStartTime, } from './sql/utils/result-formatter.js';
 export async function executeSql(args) {
     // Validate SQL input
-    if (args.sql === undefined || args.sql === null) {
-        throw new Error('sql parameter is required');
-    }
-    if (typeof args.sql !== 'string') {
-        throw new Error('sql parameter must be a string');
+    if (!args.sql || typeof args.sql !== 'string') {
+        throw new Error('sql parameter is required and must be a string');
     }
     const sql = args.sql.trim();
     if (sql.length === 0) {
@@ -189,31 +37,53 @@ export async function executeSql(args) {
     if (args.allowMultipleStatements && args.params && args.params.length > 0) {
         throw new Error('params not supported with allowMultipleStatements. Use separate execute_sql calls for parameterized queries.');
     }
+    // Connection override and transaction are mutually exclusive
+    const hasOverride = args.server || args.database || args.schema;
+    if (hasOverride && args.transactionId) {
+        throw new Error('Connection override (server/database/schema) cannot be used with transactions. Transactions are bound to the main connection.');
+    }
     const dbManager = getDbManager();
     const maxRows = validatePositiveInteger(args.maxRows, 'maxRows', 1, MAX_ROWS_LIMIT) || MAX_ROWS_DEFAULT;
     const offset = args.offset !== undefined ? validatePositiveInteger(args.offset, 'offset', 0, Number.MAX_SAFE_INTEGER) : 0;
-    // Get schema hints if requested
+    // Build connection override if specified
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
+    // Get schema hints if requested (uses current connection for schema lookup)
     let schemaHint;
     if (args.includeSchemaHint) {
-        schemaHint = await getSchemaHintForSql(sql);
+        schemaHint = await getSchemaHintForSql(sql, override);
     }
     // Handle multi-statement execution
     if (args.allowMultipleStatements) {
-        return executeMultipleStatements(sql, schemaHint, args.transactionId);
+        return executeMultipleStatements(sql, schemaHint, args.transactionId, override);
     }
     // Record start time for execution timing
-    const startTime = process.hrtime.bigint();
-    // Execute query with optional parameters (supports transaction)
+    const startTime = getStartTime();
+    // Execute query with optional parameters
     let result;
     if (args.transactionId) {
+        // Transaction-based execution (no override support)
         result = await dbManager.queryInTransaction(args.transactionId, sql, args.params);
     }
+    else if (override) {
+        // Use connection override for one-time execution
+        const { client, release, server, database, schema } = await dbManager.getClientWithOverride(override);
+        try {
+            result = await client.query(sql, args.params);
+            // Add connection info to result for transparency
+            result._connectionInfo = { server, database, schema };
+        }
+        finally {
+            release();
+        }
+    }
     else {
+        // Default: use main connection
         result = await dbManager.query(sql, args.params);
     }
-    // Calculate execution time in milliseconds
-    const endTime = process.hrtime.bigint();
-    const executionTimeMs = Number(endTime - startTime) / 1_000_000;
+    // Calculate execution time in milliseconds (already rounded to 2 decimal places)
+    const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
     // Defensive: ensure result has expected structure
     if (!result || typeof result !== 'object') {
         throw new Error('Query returned invalid result');
@@ -240,7 +110,7 @@ export async function executeSql(args) {
             offset: startIndex,
             fields,
             rows: paginatedRows,
-            executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+            executionTimeMs,
             generatedAt: new Date().toISOString()
         };
         fs.writeFileSync(filePath, JSON.stringify(outputData, null, 2), { mode: 0o600 });
@@ -250,7 +120,7 @@ export async function executeSql(args) {
             fields,
             outputFile: filePath,
             truncated: true,
-            executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+            executionTimeMs,
             offset: startIndex,
             hasMore: endIndex < totalRows,
             ...(schemaHint && { schemaHint })
@@ -260,7 +130,7 @@ export async function executeSql(args) {
         rows: paginatedRows,
         rowCount: totalRows,
         fields,
-        executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+        executionTimeMs,
         offset: startIndex,
         hasMore: endIndex < totalRows,
         ...(schemaHint && { schemaHint })
@@ -269,9 +139,9 @@ export async function executeSql(args) {
 /**
  * Execute multiple SQL statements and return results for each
  */
-async function executeMultipleStatements(sql, schemaHint, transactionId) {
+async function executeMultipleStatements(sql, schemaHint, transactionId, override) {
     const dbManager = getDbManager();
-    const startTime = process.hrtime.bigint();
+    const startTime = getStartTime();
     // Parse statements with line numbers
     const parsedStatements = splitSqlStatementsWithLineNumbers(sql);
     // Filter out empty statements and comments-only
@@ -285,42 +155,75 @@ async function executeMultipleStatements(sql, schemaHint, transactionId) {
     const results = [];
     let successCount = 0;
     let failureCount = 0;
-    for (let i = 0; i < executableStatements.length; i++) {
-        const stmt = executableStatements[i];
-        const stmtResult = {
-            statementIndex: i + 1,
-            sql: stmt.sql.length > 200 ? stmt.sql.substring(0, 200) + '...' : stmt.sql,
-            lineNumber: stmt.lineNumber,
-            success: false,
-        };
+    // For override, get a single client and execute all statements
+    if (override) {
+        const { client, release } = await dbManager.getClientWithOverride(override);
         try {
-            let result;
-            if (transactionId) {
-                result = await dbManager.queryInTransaction(transactionId, stmt.sql);
-            }
-            else {
-                result = await dbManager.query(stmt.sql);
+            for (let i = 0; i < executableStatements.length; i++) {
+                const stmt = executableStatements[i];
+                const stmtResult = {
+                    statementIndex: i + 1,
+                    sql: stmt.sql.length > SQL_TRUNCATION_SHORT ? stmt.sql.substring(0, SQL_TRUNCATION_SHORT) + '...' : stmt.sql,
+                    lineNumber: stmt.lineNumber,
+                    success: false,
+                };
+                try {
+                    const result = await client.query(stmt.sql);
+                    stmtResult.success = true;
+                    stmtResult.rows = result.rows?.slice(0, MAX_ROWS_PER_STATEMENT);
+                    stmtResult.rowCount = result.rowCount ?? result.rows?.length ?? 0;
+                    successCount++;
+                }
+                catch (error) {
+                    stmtResult.success = false;
+                    stmtResult.error = error instanceof Error ? error.message : String(error);
+                    failureCount++;
+                }
+                results.push(stmtResult);
             }
-            stmtResult.success = true;
-            stmtResult.rows = result.rows?.slice(0, 100); // Limit rows per statement
-            stmtResult.rowCount = result.rowCount ?? result.rows?.length ?? 0;
-            successCount++;
         }
-        catch (error) {
-            stmtResult.success = false;
-            stmtResult.error = error instanceof Error ? error.message : String(error);
-            failureCount++;
+        finally {
+            release();
+        }
+    }
+    else {
+        // Default execution path
+        for (let i = 0; i < executableStatements.length; i++) {
+            const stmt = executableStatements[i];
+            const stmtResult = {
+                statementIndex: i + 1,
+                sql: stmt.sql.length > SQL_TRUNCATION_SHORT ? stmt.sql.substring(0, SQL_TRUNCATION_SHORT) + '...' : stmt.sql,
+                lineNumber: stmt.lineNumber,
+                success: false,
+            };
+            try {
+                let result;
+                if (transactionId) {
+                    result = await dbManager.queryInTransaction(transactionId, stmt.sql);
+                }
+                else {
+                    result = await dbManager.query(stmt.sql);
+                }
+                stmtResult.success = true;
+                stmtResult.rows = result.rows?.slice(0, MAX_ROWS_PER_STATEMENT); // Limit rows per statement
+                stmtResult.rowCount = result.rowCount ?? result.rows?.length ?? 0;
+                successCount++;
+            }
+            catch (error) {
+                stmtResult.success = false;
+                stmtResult.error = error instanceof Error ? error.message : String(error);
+                failureCount++;
+            }
+            results.push(stmtResult);
         }
-        results.push(stmtResult);
     }
-    const endTime = process.hrtime.bigint();
-    const executionTimeMs = Number(endTime - startTime) / 1_000_000;
+    const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
     return {
         results,
         totalStatements: executableStatements.length,
         successCount,
         failureCount,
-        executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+        executionTimeMs,
         ...(schemaHint && { schemaHint })
     };
 }
@@ -340,13 +243,26 @@ export async function explainQuery(args) {
         }
     }
     const dbManager = getDbManager();
-    const client = await dbManager.getClient();
+    const hasOverride = args.server || args.database || args.schema;
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
+    // Get client - either from override or main connection
+    let clientResult = null;
+    let client;
+    if (override) {
+        clientResult = await dbManager.getClientWithOverride(override);
+        client = clientResult.client;
+    }
+    else {
+        client = await dbManager.getClient();
+    }
     try {
         // If hypothetical indexes are specified, validate and create them
         if (args.hypotheticalIndexes && args.hypotheticalIndexes.length > 0) {
             // Limit number of hypothetical indexes
-            if (args.hypotheticalIndexes.length > 10) {
-                throw new Error('Maximum 10 hypothetical indexes allowed');
+            if (args.hypotheticalIndexes.length > MAX_HYPOTHETICAL_INDEXES) {
+                throw new Error(`Maximum ${MAX_HYPOTHETICAL_INDEXES} hypothetical indexes allowed`);
             }
             // Check if hypopg extension is available
             const hypopgCheck = await client.query(`
@@ -417,7 +333,8 @@ export async function explainQuery(args) {
                 await client.query('SELECT hypopg_reset()');
             }
             catch (e) {
-                // Ignore if hypopg not available
+                // hypopg extension may not be available
+                console.debug('Could not reset hypopg:', e);
             }
         }
         if (format === 'json') {
@@ -430,39 +347,14 @@ export async function explainQuery(args) {
         };
     }
     finally {
-        client.release();
-    }
-}
-/**
- * Preprocess SQL content by removing patterns.
- * Supports both literal string matching and regex patterns.
- *
- * @param sql - The SQL content to preprocess
- * @param patterns - Array of patterns to remove from SQL content
- * @param isRegex - If true, patterns are treated as regex; if false, as literal strings
- */
-function preprocessSqlContent(sql, patterns, isRegex = false) {
-    let result = sql;
-    for (const pattern of patterns) {
-        try {
-            if (isRegex) {
-                // Treat as regex pattern (multiline by default)
-                const regex = new RegExp(pattern, 'gm');
-                result = result.replace(regex, '');
-            }
-            else {
-                // Treat as literal string - escape and match on its own line
-                const escapedPattern = pattern.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
-                const regex = new RegExp(`^\\s*${escapedPattern}\\s*$`, 'gm');
-                result = result.replace(regex, '');
-            }
+        // Release client properly based on connection type
+        if (clientResult) {
+            clientResult.release();
         }
-        catch (error) {
-            // Invalid regex - skip this pattern
-            console.error(`Warning: Invalid pattern "${pattern}": ${error instanceof Error ? error.message : String(error)}`);
+        else {
+            client.release();
         }
     }
-    return result;
 }
 /**
  * Execute a SQL file from the filesystem.
@@ -506,7 +398,12 @@ export async function executeSqlFile(args) {
     const useTransaction = args.useTransaction !== false; // Default to true
     const stopOnError = args.stopOnError !== false; // Default to true
     const validateOnly = args.validateOnly === true; // Default to false
-    const startTime = process.hrtime.bigint();
+    // Build connection override if specified
+    const hasOverride = args.server || args.database || args.schema;
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
+    const startTime = getStartTime();
     // Split SQL into statements with line number tracking
     const parsedStatements = splitSqlStatementsWithLineNumbers(sqlContent);
     const executableStatements = parsedStatements.filter(stmt => {
@@ -519,13 +416,12 @@ export async function executeSqlFile(args) {
     const totalStatements = executableStatements.length;
     // If validateOnly mode, return preview without execution
     if (validateOnly) {
-        const endTime = process.hrtime.bigint();
-        const executionTimeMs = Number(endTime - startTime) / 1_000_000;
+        const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
         // Create preview of statements
         const preview = executableStatements.map((stmt, idx) => ({
             index: idx + 1,
             lineNumber: stmt.lineNumber,
-            sql: stmt.sql.length > 300 ? stmt.sql.substring(0, 300) + '...' : stmt.sql,
+            sql: stmt.sql.length > SQL_TRUNCATION_LONG ? stmt.sql.substring(0, SQL_TRUNCATION_LONG) + '...' : stmt.sql,
             type: detectStatementType(stmt.sql)
         }));
         return {
@@ -535,13 +431,22 @@ export async function executeSqlFile(args) {
             totalStatements,
             statementsExecuted: 0,
             statementsFailed: 0,
-            executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+            executionTimeMs,
             rowsAffected: 0,
             validateOnly: true,
             preview
         };
     }
-    const client = await dbManager.getClient();
+    // Get client - either from override or main connection
+    let clientResult = null;
+    let client;
+    if (override) {
+        clientResult = await dbManager.getClientWithOverride(override);
+        client = clientResult.client;
+    }
+    else {
+        client = await dbManager.getClient();
+    }
     let statementsExecuted = 0;
     let statementsFailed = 0;
     let totalRowsAffected = 0;
@@ -573,7 +478,7 @@ export async function executeSqlFile(args) {
                     collectedErrors.push({
                         statementIndex: statementIndex + 1,
                         lineNumber: stmt.lineNumber,
-                        sql: trimmed.length > 200 ? trimmed.substring(0, 200) + '...' : trimmed,
+                        sql: trimmed.length > SQL_TRUNCATION_SHORT ? trimmed.substring(0, SQL_TRUNCATION_SHORT) + '...' : trimmed,
                         error: errorMessage
                     });
                     throw error;
@@ -582,7 +487,7 @@ export async function executeSqlFile(args) {
                 collectedErrors.push({
                     statementIndex: statementIndex + 1,
                     lineNumber: stmt.lineNumber,
-                    sql: trimmed.length > 200 ? trimmed.substring(0, 200) + '...' : trimmed,
+                    sql: trimmed.length > SQL_TRUNCATION_SHORT ? trimmed.substring(0, SQL_TRUNCATION_SHORT) + '...' : trimmed,
                     error: errorMessage
                 });
                 console.error(`Warning: Statement ${statementIndex + 1} at line ${stmt.lineNumber} failed: ${errorMessage}`);
@@ -591,8 +496,7 @@ export async function executeSqlFile(args) {
         if (useTransaction && !rolledBack) {
             await client.query('COMMIT');
         }
-        const endTime = process.hrtime.bigint();
-        const executionTimeMs = Number(endTime - startTime) / 1_000_000;
+        const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
         const result = {
             success: statementsFailed === 0,
             filePath: resolvedPath,
@@ -600,7 +504,7 @@ export async function executeSqlFile(args) {
             totalStatements,
             statementsExecuted,
             statementsFailed,
-            executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+            executionTimeMs,
             rowsAffected: totalRowsAffected
         };
         // Include errors array if there were any failures (when stopOnError=false)
@@ -610,8 +514,7 @@ export async function executeSqlFile(args) {
         return result;
     }
     catch (error) {
-        const endTime = process.hrtime.bigint();
-        const executionTimeMs = Number(endTime - startTime) / 1_000_000;
+        const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
         const result = {
             success: false,
             filePath: resolvedPath,
@@ -619,7 +522,7 @@ export async function executeSqlFile(args) {
             totalStatements,
             statementsExecuted,
             statementsFailed,
-            executionTimeMs: Math.round(executionTimeMs * 100) / 100,
+            executionTimeMs,
             rowsAffected: totalRowsAffected,
             error: error instanceof Error ? error.message : String(error),
             rollback: rolledBack
@@ -630,40 +533,14 @@ export async function executeSqlFile(args) {
         return result;
     }
     finally {
-        client.release();
-    }
-}
-/**
- * Detect the type of SQL statement (SELECT, INSERT, UPDATE, DELETE, CREATE, etc.)
- */
-function detectStatementType(sql) {
-    const trimmed = stripLeadingComments(sql).toUpperCase();
-    // Common statement types
-    const types = [
-        'SELECT', 'INSERT', 'UPDATE', 'DELETE', 'CREATE', 'ALTER', 'DROP',
-        'TRUNCATE', 'GRANT', 'REVOKE', 'BEGIN', 'COMMIT', 'ROLLBACK',
-        'SET', 'SHOW', 'EXPLAIN', 'ANALYZE', 'VACUUM', 'REINDEX',
-        'COMMENT', 'WITH', 'DO', 'CALL', 'EXECUTE'
-    ];
-    for (const type of types) {
-        if (trimmed.startsWith(type + ' ') || trimmed.startsWith(type + '\n') ||
-            trimmed.startsWith(type + '\t') || trimmed === type) {
-            // Special case for WITH - check if it's a CTE followed by SELECT/INSERT/UPDATE/DELETE
-            if (type === 'WITH') {
-                if (trimmed.includes('SELECT'))
-                    return 'WITH SELECT';
-                if (trimmed.includes('INSERT'))
-                    return 'WITH INSERT';
-                if (trimmed.includes('UPDATE'))
-                    return 'WITH UPDATE';
-                if (trimmed.includes('DELETE'))
-                    return 'WITH DELETE';
-                return 'WITH';
-            }
-            return type;
+        // Release client properly based on connection type
+        if (clientResult) {
+            clientResult.release();
+        }
+        else {
+            client.release();
         }
     }
-    return 'UNKNOWN';
 }
 /**
  * Preview a SQL file without executing.
@@ -703,7 +580,7 @@ export async function previewSqlFile(args) {
     if (args.stripPatterns && args.stripPatterns.length > 0) {
         sqlContent = preprocessSqlContent(sqlContent, args.stripPatterns, args.stripAsRegex === true);
     }
-    const maxStatements = Math.min(args.maxStatements || 20, 100);
+    const maxStatements = Math.min(args.maxStatements || DEFAULT_PREVIEW_STATEMENTS, MAX_PREVIEW_STATEMENTS);
     // Split SQL into statements with line number tracking
     const parsedStatements = splitSqlStatementsWithLineNumbers(sqlContent);
     const executableStatements = parsedStatements.filter(stmt => {
@@ -739,15 +616,20 @@ export async function previewSqlFile(args) {
     const statements = executableStatements.slice(0, maxStatements).map((stmt, idx) => ({
         index: idx + 1,
         lineNumber: stmt.lineNumber,
-        sql: stmt.sql.length > 300 ? stmt.sql.substring(0, 300) + '...' : stmt.sql,
+        sql: stmt.sql.length > SQL_TRUNCATION_LONG ? stmt.sql.substring(0, SQL_TRUNCATION_LONG) + '...' : stmt.sql,
         type: detectStatementType(stmt.sql)
     }));
     // Format file size
-    const fileSizeFormatted = stats.size < 1024
-        ? `${stats.size} bytes`
-        : stats.size < 1024 * 1024
-            ? `${(stats.size / 1024).toFixed(1)} KB`
-            : `${(stats.size / (1024 * 1024)).toFixed(2)} MB`;
+    let fileSizeFormatted;
+    if (stats.size < 1024) {
+        fileSizeFormatted = `${stats.size} bytes`;
+    }
+    else if (stats.size < 1024 * 1024) {
+        fileSizeFormatted = `${(stats.size / 1024).toFixed(1)} KB`;
+    }
+    else {
+        fileSizeFormatted = `${(stats.size / (1024 * 1024)).toFixed(2)} MB`;
+    }
     // Generate summary
     const typeEntries = Object.entries(statementsByType).sort((a, b) => b[1] - a[1]);
     const typeSummary = typeEntries.map(([type, count]) => `${count} ${type}`).join(', ');
@@ -808,7 +690,7 @@ export async function dryRunSqlFile(args) {
     if (args.stripPatterns && args.stripPatterns.length > 0) {
         sqlContent = preprocessSqlContent(sqlContent, args.stripPatterns, args.stripAsRegex === true);
     }
-    const maxStatements = Math.min(args.maxStatements || 50, 200);
+    const maxStatements = Math.min(args.maxStatements || DEFAULT_DRY_RUN_STATEMENTS, MAX_DRY_RUN_STATEMENTS);
     const stopOnError = args.stopOnError === true;
     // Split SQL into statements with line number tracking
     const parsedStatements = splitSqlStatementsWithLineNumbers(sqlContent);
@@ -827,14 +709,33 @@ export async function dryRunSqlFile(args) {
         nonRollbackableWarnings.push(...warnings);
     });
     // Format file size
-    const fileSizeFormatted = stats.size < 1024
-        ? `${stats.size} bytes`
-        : stats.size < 1024 * 1024
-            ? `${(stats.size / 1024).toFixed(1)} KB`
-            : `${(stats.size / (1024 * 1024)).toFixed(2)} MB`;
+    let fileSizeFormatted;
+    if (stats.size < 1024) {
+        fileSizeFormatted = `${stats.size} bytes`;
+    }
+    else if (stats.size < 1024 * 1024) {
+        fileSizeFormatted = `${(stats.size / 1024).toFixed(1)} KB`;
+    }
+    else {
+        fileSizeFormatted = `${(stats.size / (1024 * 1024)).toFixed(2)} MB`;
+    }
     const dbManager = getDbManager();
-    const client = await dbManager.getClient();
-    const startTime = process.hrtime.bigint();
+    // Build connection override if specified
+    const hasOverride = args.server || args.database || args.schema;
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
+    // Get client - either from override or main connection
+    let clientResult = null;
+    let client;
+    if (override) {
+        clientResult = await dbManager.getClientWithOverride(override);
+        client = clientResult.client;
+    }
+    else {
+        client = await dbManager.getClient();
+    }
+    const startTime = getStartTime();
     const statementResults = [];
     const statementsByType = {};
     let successCount = 0;
@@ -849,11 +750,11 @@ export async function dryRunSqlFile(args) {
             const stmt = executableStatements[idx];
             const stmtType = detectStatementType(stmt.sql);
             statementsByType[stmtType] = (statementsByType[stmtType] || 0) + 1;
-            const stmtStartTime = process.hrtime.bigint();
+            const stmtStartTime = getStartTime();
             const result = {
                 index: idx + 1,
                 lineNumber: stmt.lineNumber,
-                sql: stmt.sql.length > 300 ? stmt.sql.substring(0, 300) + '...' : stmt.sql,
+                sql: stmt.sql.length > SQL_TRUNCATION_LONG ? stmt.sql.substring(0, SQL_TRUNCATION_LONG) + '...' : stmt.sql,
                 type: stmtType,
                 success: false
             };
@@ -908,8 +809,7 @@ export async function dryRunSqlFile(args) {
                     }
                 }
             }
-            const stmtEndTime = process.hrtime.bigint();
-            result.executionTimeMs = Math.round(Number(stmtEndTime - stmtStartTime) / 1_000_000 * 100) / 100;
+            result.executionTimeMs = calculateExecutionTime(stmtStartTime, getStartTime());
             // Only include results up to maxStatements
             if (statementResults.length < maxStatements) {
                 statementResults.push(result);
@@ -929,10 +829,15 @@ export async function dryRunSqlFile(args) {
         throw e;
     }
     finally {
-        client.release();
+        // Release client properly based on connection type
+        if (clientResult) {
+            clientResult.release();
+        }
+        else {
+            client.release();
+        }
     }
-    const endTime = process.hrtime.bigint();
-    const executionTimeMs = Math.round(Number(endTime - startTime) / 1_000_000 * 100) / 100;
+    const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
     // Generate summary
     const typeEntries = Object.entries(statementsByType).sort((a, b) => b[1] - a[1]);
     const typeSummary = typeEntries.map(([type, count]) => `${count} ${type}`).join(', ');
@@ -966,194 +871,99 @@ export async function dryRunSqlFile(args) {
     };
 }
 /**
- * Strips leading line comments and block comments from SQL to check if there's actual SQL.
- * Returns empty string if the entire content is just comments.
- */
-function stripLeadingComments(sql) {
-    let result = sql.trim();
-    while (result.length > 0) {
-        // Strip leading line comments
-        if (result.startsWith('--')) {
-            const newlineIndex = result.indexOf('\n');
-            if (newlineIndex === -1) {
-                return ''; // Entire string is a line comment
-            }
-            result = result.substring(newlineIndex + 1).trim();
-            continue;
-        }
-        // Strip leading block comments
-        if (result.startsWith('/*')) {
-            const endIndex = result.indexOf('*/');
-            if (endIndex === -1) {
-                return ''; // Unclosed block comment
-            }
-            result = result.substring(endIndex + 2).trim();
-            continue;
-        }
-        // No more leading comments
-        break;
-    }
-    return result;
-}
-/**
- * Split SQL content into individual statements with line number tracking.
- * Returns ParsedStatement objects with SQL and line number info.
+ * Gets schema hints for tables mentioned in SQL
  */
-function splitSqlStatementsWithLineNumbers(sql) {
-    const statements = [];
-    let current = '';
-    let currentLineNumber = 1;
-    let statementStartLine = 1;
-    let inString = false;
-    let stringChar = '';
-    let inLineComment = false;
-    let inBlockComment = false;
-    let i = 0;
-    while (i < sql.length) {
-        const char = sql[i];
-        const nextChar = sql[i + 1] || '';
-        // Track line numbers
-        if (char === '\n') {
-            currentLineNumber++;
-        }
-        // If starting a new statement (current is empty/whitespace), record line number
-        if (current.trim() === '' && char.trim() !== '') {
-            statementStartLine = currentLineNumber;
-        }
-        // Handle line comments
-        if (!inString && !inBlockComment && char === '-' && nextChar === '-') {
-            inLineComment = true;
-            current += char;
-            i++;
-            continue;
-        }
-        if (inLineComment && (char === '\n' || char === '\r')) {
-            inLineComment = false;
-            current += char;
-            i++;
-            continue;
-        }
-        // Handle block comments
-        if (!inString && !inLineComment && char === '/' && nextChar === '*') {
-            inBlockComment = true;
-            current += char + nextChar;
-            i += 2;
-            continue;
-        }
-        if (inBlockComment && char === '*' && nextChar === '/') {
-            inBlockComment = false;
-            current += char + nextChar;
-            i += 2;
-            continue;
-        }
-        // Handle string literals
-        if (!inLineComment && !inBlockComment && (char === "'" || char === '"')) {
-            if (!inString) {
-                inString = true;
-                stringChar = char;
-            }
-            else if (char === stringChar) {
-                if (nextChar === stringChar) {
-                    current += char + nextChar;
-                    i += 2;
-                    continue;
+async function getSchemaHintForSql(sql, override) {
+    const dbManager = getDbManager();
+    const tables = extractTablesFromSql(sql);
+    const tableHints = [];
+    // If override is specified, use a single client for all queries
+    if (override) {
+        const { client, release } = await dbManager.getClientWithOverride(override);
+        try {
+            for (const { schema, table } of tables.slice(0, MAX_TABLES_TO_ANALYZE)) {
+                try {
+                    const columnsResult = await client.query(`
+            SELECT
+              column_name as name,
+              data_type as type,
+              is_nullable = 'YES' as nullable
+            FROM information_schema.columns
+            WHERE table_schema = $1 AND table_name = $2
+            ORDER BY ordinal_position
+          `, [schema, table]);
+                    const pkResult = await client.query(`
+            SELECT a.attname as column_name
+            FROM pg_index i
+            JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
+            JOIN pg_class c ON c.oid = i.indrelid
+            JOIN pg_namespace n ON n.oid = c.relnamespace
+            WHERE i.indisprimary
+              AND n.nspname = $1
+              AND c.relname = $2
+          `, [schema, table]);
+                    const fkResult = await client.query(`
+            SELECT
+              kcu.column_name,
+              ccu.table_schema || '.' || ccu.table_name as referenced_table,
+              ccu.column_name as referenced_column
+            FROM information_schema.table_constraints tc
+            JOIN information_schema.key_column_usage kcu
+              ON tc.constraint_name = kcu.constraint_name
+              AND tc.table_schema = kcu.table_schema
+            JOIN information_schema.constraint_column_usage ccu
+              ON ccu.constraint_name = tc.constraint_name
+            WHERE tc.constraint_type = 'FOREIGN KEY'
+              AND tc.table_schema = $1
+              AND tc.table_name = $2
+          `, [schema, table]);
+                    const countResult = await client.query(`
+            SELECT reltuples::bigint as estimate
+            FROM pg_class c
+            JOIN pg_namespace n ON n.oid = c.relnamespace
+            WHERE n.nspname = $1 AND c.relname = $2
+          `, [schema, table]);
+                    const hint = {
+                        schema,
+                        table,
+                        columns: columnsResult.rows.map(r => ({
+                            name: r.name,
+                            type: r.type,
+                            nullable: r.nullable
+                        })),
+                        primaryKey: pkResult.rows.map(r => r.column_name),
+                        rowCountEstimate: countResult.rows[0]?.estimate || 0
+                    };
+                    if (fkResult.rows.length > 0) {
+                        const fkMap = new Map();
+                        for (const row of fkResult.rows) {
+                            const key = row.referenced_table;
+                            if (!fkMap.has(key)) {
+                                fkMap.set(key, { columns: [], referencedColumns: [] });
+                            }
+                            fkMap.get(key).columns.push(row.column_name);
+                            fkMap.get(key).referencedColumns.push(row.referenced_column);
+                        }
+                        hint.foreignKeys = Array.from(fkMap.entries()).map(([refTable, data]) => ({
+                            columns: data.columns,
+                            referencedTable: refTable,
+                            referencedColumns: data.referencedColumns
+                        }));
+                    }
+                    tableHints.push(hint);
                 }
-                inString = false;
-                stringChar = '';
-            }
-        }
-        // Handle dollar-quoted strings (PostgreSQL specific)
-        if (!inString && !inLineComment && !inBlockComment && char === '$') {
-            const dollarMatch = sql.slice(i).match(/^(\$[a-zA-Z0-9_]*\$)/);
-            if (dollarMatch) {
-                const dollarTag = dollarMatch[1];
-                const endIndex = sql.indexOf(dollarTag, i + dollarTag.length);
-                if (endIndex !== -1) {
-                    const dollarContent = sql.slice(i, endIndex + dollarTag.length);
-                    // Count newlines in dollar-quoted content
-                    const newlines = (dollarContent.match(/\n/g) || []).length;
-                    currentLineNumber += newlines;
-                    current += dollarContent;
-                    i = endIndex + dollarTag.length;
-                    continue;
+                catch (error) {
+                    console.error(`Could not get schema hint for ${schema}.${table}: ${error}`);
                 }
             }
         }
-        // Handle statement separator
-        if (!inString && !inLineComment && !inBlockComment && char === ';') {
-            current += char;
-            const trimmed = current.trim();
-            if (trimmed) {
-                statements.push({ sql: trimmed, lineNumber: statementStartLine });
-            }
-            current = '';
-            statementStartLine = currentLineNumber;
-            i++;
-            continue;
-        }
-        current += char;
-        i++;
-    }
-    // Add remaining content if any
-    const trimmed = current.trim();
-    if (trimmed) {
-        statements.push({ sql: trimmed, lineNumber: statementStartLine });
-    }
-    return statements;
-}
-/**
- * Extracts table names from a SQL query.
- * Handles common patterns: FROM, JOIN, INTO, UPDATE, DELETE FROM
- */
-function extractTablesFromSql(sql) {
-    const tables = [];
-    const seen = new Set();
-    // Normalize SQL: remove comments and extra whitespace
-    const normalized = sql
-        .replace(/--[^\n]*/g, '') // Remove line comments
-        .replace(/\/\*[\s\S]*?\*\//g, '') // Remove block comments
-        .replace(/\s+/g, ' ') // Normalize whitespace
-        .trim();
-    // Patterns to find table references
-    const patterns = [
-        /\bFROM\s+(["`]?[\w]+["`]?(?:\s*\.\s*["`]?[\w]+["`]?)?)/gi,
-        /\bJOIN\s+(["`]?[\w]+["`]?(?:\s*\.\s*["`]?[\w]+["`]?)?)/gi,
-        /\bINTO\s+(["`]?[\w]+["`]?(?:\s*\.\s*["`]?[\w]+["`]?)?)/gi,
-        /\bUPDATE\s+(["`]?[\w]+["`]?(?:\s*\.\s*["`]?[\w]+["`]?)?)/gi,
-        /\bDELETE\s+FROM\s+(["`]?[\w]+["`]?(?:\s*\.\s*["`]?[\w]+["`]?)?)/gi,
-    ];
-    for (const pattern of patterns) {
-        let match;
-        while ((match = pattern.exec(normalized)) !== null) {
-            const tableRef = match[1].replace(/["`]/g, '').trim();
-            // Skip common SQL keywords that might be matched
-            if (['SELECT', 'WHERE', 'SET', 'VALUES', 'AND', 'OR'].includes(tableRef.toUpperCase())) {
-                continue;
-            }
-            let schema = 'public';
-            let table = tableRef;
-            if (tableRef.includes('.')) {
-                const parts = tableRef.split('.');
-                schema = parts[0].trim();
-                table = parts[1].trim();
-            }
-            const key = `${schema}.${table}`.toLowerCase();
-            if (!seen.has(key)) {
-                seen.add(key);
-                tables.push({ schema, table });
-            }
+        finally {
+            release();
         }
+        return { tables: tableHints };
     }
-    return tables;
-}
-/**
- * Gets schema hints for tables mentioned in SQL
- */
-async function getSchemaHintForSql(sql) {
-    const dbManager = getDbManager();
-    const tables = extractTablesFromSql(sql);
-    const tableHints = [];
-    for (const { schema, table } of tables.slice(0, 10)) { // Limit to 10 tables
+    // Default path: use main connection
+    for (const { schema, table } of tables.slice(0, MAX_TABLES_TO_ANALYZE)) { // Limit to 10 tables
         try {
             // Get columns
             const columnsResult = await dbManager.query(`
@@ -1245,10 +1055,10 @@ export async function mutationPreview(args) {
         throw new Error('sql parameter is required');
     }
     const sql = args.sql.trim();
-    const sampleSize = Math.min(args.sampleSize || 5, 20); // Default 5, max 20
+    const sampleSize = Math.min(args.sampleSize || DEFAULT_MUTATION_SAMPLE_SIZE, MAX_MUTATION_SAMPLE_SIZE);
     // Detect mutation type
     const upperSql = sql.toUpperCase();
-    let mutationType = 'UNKNOWN';
+    let mutationType;
     if (upperSql.startsWith('UPDATE')) {
         mutationType = 'UPDATE';
     }
@@ -1262,10 +1072,15 @@ export async function mutationPreview(args) {
         throw new Error('SQL must be an INSERT, UPDATE, or DELETE statement');
     }
     const dbManager = getDbManager();
+    // Build connection override if specified
+    const hasOverride = args.server || args.database || args.schema;
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
     // For INSERT, we can't preview affected rows
     if (mutationType === 'INSERT') {
         // Use EXPLAIN to estimate rows
-        const explainResult = await dbManager.query(`EXPLAIN (FORMAT JSON) ${sql}`);
+        const explainResult = await dbManager.queryWithOverride(`EXPLAIN (FORMAT JSON) ${sql}`, undefined, override);
         const plan = explainResult.rows[0]['QUERY PLAN'][0];
         return {
             mutationType,
@@ -1277,19 +1092,22 @@ export async function mutationPreview(args) {
     // For UPDATE and DELETE, extract WHERE clause and table
     let targetTable;
     let whereClause;
+    // Extract WHERE clause using string search (avoids regex backtracking)
+    const whereIdx = sql.toUpperCase().lastIndexOf('WHERE');
+    if (whereIdx !== -1) {
+        whereClause = sql.substring(whereIdx + 5).trim();
+    }
+    const updatePattern = /UPDATE\s+(["`]?\w[\w.]*["`]?)\s+SET/i;
+    const deletePattern = /DELETE\s+FROM\s+(["`]?\w[\w.]*["`]?)/i;
     if (mutationType === 'UPDATE') {
         // Pattern: UPDATE table SET ... WHERE ...
-        const updateMatch = sql.match(/UPDATE\s+(["`]?[\w.]+["`]?)\s+SET/i);
-        const whereMatch = sql.match(/\bWHERE\s+(.+)$/is);
+        const updateMatch = updatePattern.exec(sql);
         targetTable = updateMatch?.[1]?.replace(/["`]/g, '');
-        whereClause = whereMatch?.[1];
     }
     else if (mutationType === 'DELETE') {
         // Pattern: DELETE FROM table WHERE ...
-        const deleteMatch = sql.match(/DELETE\s+FROM\s+(["`]?[\w.]+["`]?)/i);
-        const whereMatch = sql.match(/\bWHERE\s+(.+)$/is);
+        const deleteMatch = deletePattern.exec(sql);
         targetTable = deleteMatch?.[1]?.replace(/["`]/g, '');
-        whereClause = whereMatch?.[1];
     }
     if (!targetTable) {
         throw new Error('Could not parse target table from SQL');
@@ -1297,12 +1115,13 @@ export async function mutationPreview(args) {
     // Get estimated row count using EXPLAIN
     let estimatedRowsAffected = 0;
     try {
-        const explainResult = await dbManager.query(`EXPLAIN (FORMAT JSON) ${sql}`);
+        const explainResult = await dbManager.queryWithOverride(`EXPLAIN (FORMAT JSON) ${sql}`, undefined, override);
         const plan = explainResult.rows[0]['QUERY PLAN'][0];
         estimatedRowsAffected = plan?.Plan?.['Plan Rows'] || 0;
     }
     catch (error) {
-        // EXPLAIN might fail, continue with count query
+        // EXPLAIN might fail for complex queries, continue with count query
+        console.debug('Could not get EXPLAIN plan:', error);
     }
     // Build SELECT query to get sample of affected rows
     let sampleRows = [];
@@ -1310,14 +1129,14 @@ export async function mutationPreview(args) {
         const selectSql = whereClause
             ? `SELECT * FROM ${targetTable} WHERE ${whereClause} LIMIT ${sampleSize}`
             : `SELECT * FROM ${targetTable} LIMIT ${sampleSize}`;
-        const sampleResult = await dbManager.query(selectSql);
+        const sampleResult = await dbManager.queryWithOverride(selectSql, undefined, override);
         sampleRows = sampleResult.rows;
         // If EXPLAIN didn't work, get count
         if (estimatedRowsAffected === 0) {
             const countSql = whereClause
                 ? `SELECT COUNT(*) as cnt FROM ${targetTable} WHERE ${whereClause}`
                 : `SELECT COUNT(*) as cnt FROM ${targetTable}`;
-            const countResult = await dbManager.query(countSql);
+            const countResult = await dbManager.queryWithOverride(countSql, undefined, override);
             estimatedRowsAffected = parseInt(countResult.rows[0]?.cnt || '0', 10);
         }
     }
@@ -1353,17 +1172,20 @@ export async function mutationDryRun(args) {
         throw new Error('sql parameter is required');
     }
     const sql = args.sql.trim();
-    const sampleSize = Math.min(args.sampleSize || MAX_DRY_RUN_SAMPLE_ROWS, 20);
+    const sampleSize = Math.min(args.sampleSize || MAX_DRY_RUN_SAMPLE_ROWS, MAX_MUTATION_SAMPLE_SIZE);
     // Detect mutation type
     const upperSql = sql.toUpperCase();
-    let mutationType = 'UNKNOWN';
-    if (upperSql.startsWith('UPDATE') || upperSql.match(/^WITH\b.*\bUPDATE\b/s)) {
+    let mutationType;
+    const withUpdatePattern = /^WITH\b.*\bUPDATE\b/s;
+    const withDeletePattern = /^WITH\b.*\bDELETE\b/s;
+    const withInsertPattern = /^WITH\b.*\bINSERT\b/s;
+    if (upperSql.startsWith('UPDATE') || withUpdatePattern.test(upperSql)) {
         mutationType = 'UPDATE';
     }
-    else if (upperSql.startsWith('DELETE') || upperSql.match(/^WITH\b.*\bDELETE\b/s)) {
+    else if (upperSql.startsWith('DELETE') || withDeletePattern.test(upperSql)) {
         mutationType = 'DELETE';
     }
-    else if (upperSql.startsWith('INSERT') || upperSql.match(/^WITH\b.*\bINSERT\b/s)) {
+    else if (upperSql.startsWith('INSERT') || withInsertPattern.test(upperSql)) {
         mutationType = 'INSERT';
     }
     else {
@@ -1372,14 +1194,19 @@ export async function mutationDryRun(args) {
     // Check for non-rollbackable operations
     const nonRollbackableWarnings = detectNonRollbackableOperations(sql);
     const mustSkipWarnings = nonRollbackableWarnings.filter(w => w.mustSkip);
+    const dbManager = getDbManager();
+    // Build connection override if specified
+    const hasOverride = args.server || args.database || args.schema;
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
     // Skip only if there are mustSkip warnings
     if (mustSkipWarnings.length > 0) {
         const skipReason = mustSkipWarnings.map(w => w.message).join('; ');
         // Run EXPLAIN to show query plan without executing
         let explainPlan;
-        const dbManager = getDbManager();
         try {
-            const explainResult = await dbManager.query(`EXPLAIN (FORMAT JSON) ${sql}`);
+            const explainResult = await dbManager.queryWithOverride(`EXPLAIN (FORMAT JSON) ${sql}`, undefined, override);
             if (explainResult.rows && explainResult.rows.length > 0) {
                 explainPlan = explainResult.rows[0]['QUERY PLAN'];
             }
@@ -1401,25 +1228,41 @@ export async function mutationDryRun(args) {
     // Extract table and WHERE clause
     let targetTable;
     let whereClause;
+    // Extract WHERE clause using string search (avoids regex backtracking)
+    const upperSqlForWhere = sql.toUpperCase();
+    const whereIdx = upperSqlForWhere.lastIndexOf('WHERE');
+    if (whereIdx !== -1) {
+        let endIdx = upperSqlForWhere.indexOf('RETURNING', whereIdx);
+        if (endIdx === -1)
+            endIdx = sql.length;
+        whereClause = sql.substring(whereIdx + 5, endIdx).trim();
+    }
+    const updateTablePattern = /UPDATE\s+(["`]?\w[\w.]*["`]?)\s+SET/i;
+    const deleteTablePattern = /DELETE\s+FROM\s+(["`]?\w[\w.]*["`]?)/i;
+    const insertTablePattern = /INSERT\s+INTO\s+(["`]?\w[\w.]*["`]?)/i;
     if (mutationType === 'UPDATE') {
-        const updateMatch = sql.match(/UPDATE\s+(["`]?[\w.]+["`]?)\s+SET/i);
-        const whereMatch = sql.match(/\bWHERE\s+(.+?)(?:RETURNING|$)/is);
+        const updateMatch = updateTablePattern.exec(sql);
         targetTable = updateMatch?.[1]?.replace(/["`]/g, '');
-        whereClause = whereMatch?.[1]?.trim();
     }
     else if (mutationType === 'DELETE') {
-        const deleteMatch = sql.match(/DELETE\s+FROM\s+(["`]?[\w.]+["`]?)/i);
-        const whereMatch = sql.match(/\bWHERE\s+(.+?)(?:RETURNING|$)/is);
+        const deleteMatch = deleteTablePattern.exec(sql);
         targetTable = deleteMatch?.[1]?.replace(/["`]/g, '');
-        whereClause = whereMatch?.[1]?.trim();
     }
     else if (mutationType === 'INSERT') {
-        const insertMatch = sql.match(/INSERT\s+INTO\s+(["`]?[\w.]+["`]?)/i);
+        const insertMatch = insertTablePattern.exec(sql);
         targetTable = insertMatch?.[1]?.replace(/["`]/g, '');
     }
-    const dbManager = getDbManager();
-    const client = await dbManager.getClient();
-    const startTime = process.hrtime.bigint();
+    // Get client - either from override or main connection
+    let clientResult = null;
+    let client;
+    if (override) {
+        clientResult = await dbManager.getClientWithOverride(override);
+        client = clientResult.client;
+    }
+    else {
+        client = await dbManager.getClient();
+    }
+    const startTime = getStartTime();
     let beforeRows;
     let affectedRows = [];
     let rowsAffected = 0;
@@ -1494,15 +1337,20 @@ export async function mutationDryRun(args) {
         error = extractDryRunError(e);
     }
     finally {
-        client.release();
+        // Release client properly based on connection type
+        if (clientResult) {
+            clientResult.release();
+        }
+        else {
+            client.release();
+        }
     }
-    const endTime = process.hrtime.bigint();
-    const executionTimeMs = Number(endTime - startTime) / 1_000_000;
+    const executionTimeMs = calculateExecutionTime(startTime, getStartTime());
     const result = {
         mutationType,
         success,
         rowsAffected,
-        executionTimeMs: Math.round(executionTimeMs * 100) / 100
+        executionTimeMs
     };
     if (beforeRows && beforeRows.length > 0) {
         result.beforeRows = beforeRows;
@@ -1541,8 +1389,8 @@ export async function batchExecute(args) {
     if (args.queries.length === 0) {
         throw new Error('queries array cannot be empty');
     }
-    if (args.queries.length > 20) {
-        throw new Error('Maximum 20 queries allowed in a batch');
+    if (args.queries.length > MAX_BATCH_QUERIES) {
+        throw new Error(`Maximum ${MAX_BATCH_QUERIES} queries allowed in a batch`);
     }
     // Validate each query
     const seenNames = new Set();
@@ -1560,36 +1408,39 @@ export async function batchExecute(args) {
     }
     const dbManager = getDbManager();
     const stopOnError = args.stopOnError === true; // Default false
-    const startTime = process.hrtime.bigint();
+    const startTime = getStartTime();
+    // Build connection override if specified
+    const hasOverride = args.server || args.database || args.schema;
+    const override = hasOverride
+        ? { server: args.server, database: args.database, schema: args.schema }
+        : undefined;
     const results = {};
     let successCount = 0;
     let failureCount = 0;
     // Execute all queries in parallel
     const promises = args.queries.map(async (query) => {
-        const queryStartTime = process.hrtime.bigint();
+        const queryStartTime = getStartTime();
         try {
-            const result = await dbManager.query(query.sql, query.params);
-            const queryEndTime = process.hrtime.bigint();
-            const executionTimeMs = Number(queryEndTime - queryStartTime) / 1_000_000;
+            const result = await dbManager.queryWithOverride(query.sql, query.params, override);
+            const executionTimeMs = calculateExecutionTime(queryStartTime, getStartTime());
             return {
                 name: query.name,
                 result: {
                     success: true,
                     rows: result.rows,
                     rowCount: result.rowCount ?? result.rows.length,
-                    executionTimeMs: Math.round(executionTimeMs * 100) / 100
+                    executionTimeMs
                 }
             };
         }
         catch (error) {
-            const queryEndTime = process.hrtime.bigint();
-            const executionTimeMs = Number(queryEndTime - queryStartTime) / 1_000_000;
+            const executionTimeMs = calculateExecutionTime(queryStartTime, getStartTime());
             return {
                 name: query.name,
                 result: {
                     success: false,
                     error: error instanceof Error ? error.message : String(error),
-                    executionTimeMs: Math.round(executionTimeMs * 100) / 100
+                    executionTimeMs
                 }
             };
         }
@@ -1610,13 +1461,12 @@ export async function batchExecute(args) {
             }
         }
     }
-    const endTime = process.hrtime.bigint();
-    const totalExecutionTimeMs = Number(endTime - startTime) / 1_000_000;
+    const totalExecutionTimeMs = calculateExecutionTime(startTime, getStartTime());
     return {
         totalQueries: args.queries.length,
         successCount,
         failureCount,
-        totalExecutionTimeMs: Math.round(totalExecutionTimeMs * 100) / 100,
+        totalExecutionTimeMs,
         results
     };
 }
@@ -1626,11 +1476,12 @@ export async function batchExecute(args) {
 export async function beginTransaction(args) {
     const dbManager = getDbManager();
     const info = await dbManager.beginTransaction(args?.name);
+    const namePart = info.name ? ` "${info.name}"` : '';
     return {
         transactionId: info.transactionId,
         name: info.name,
         status: 'started',
-        message: `Transaction${info.name ? ` "${info.name}"` : ''} started. Use transactionId "${info.transactionId}" with execute_sql or commit/rollback.`
+        message: `Transaction${namePart} started. Use transactionId "${info.transactionId}" with execute_sql or commit/rollback.`
     };
 }
 /**