feat: add failed query retries

Author: naorpeled

README.md

@@ -249,6 +249,9 @@ Below is a table containing all of the possible configuration options for `serve
 | zombieMaxTimeout | `Integer` | The maximum number of seconds that a connection can stay idle before being recycled. | `900` |
 | zombieMinTimeout | `Integer` | The minimum number of *seconds* that a connection must be idle before the module will recycle it. | `3` |
 | returnFinalSqlQuery | `Boolean` | Flag indicating whether to attach the final SQL query (with substituted values) to the results. When enabled, the SQL query will be available as a non-enumerable `sql` property on array results or as a regular property on object results. | `false` |
+| maxQueryRetries | `Integer` | Maximum number of times to retry a query before giving up. | `0` |
+| queryRetryBackoff | `String` or `Function` | Backoff algorithm to be used when retrying queries. Possible values are `full` and `decorrelated`, or you can also specify your own algorithm. See [Connection Backoff](#connection-backoff) for more information. | `full` |
+| onQueryRetry | `function` | [Event](#events) callback when queries are retried. | |
 
 ### Connection Backoff
 If `manageConns` is not set to `false`, then this module will automatically kill idle connections or disconnect the current connection if the `connUtilization` limit is reached. Even with this aggressive strategy, it is possible that multiple functions will be competing for available connections. The `backoff` setting uses the strategy outlined [here](https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/) to use *Jitter* instead of *Exponential Backoff* when attempting connection retries.
@@ -323,6 +326,12 @@ The `onConnect` event recieves the MySQL `connection` object, `onKill` receives
 onRetry: (err,retries,delay,type) => { console.log('RETRY') }
 ```
 
+`onQueryRetry` also receives *four* arguments. The `error` object, the number of `retries`, the `delay` until the next retry, and the `backoff` algorithm used (`full`, `decorrelated` or `custom`).
+
+```javascript
+onQueryRetry: (err,retries,delay,type) => { console.log('QUERY RETRY') }
+```
+
 ## MySQL Server Configuration
 There really isn't anything special that needs to be done in order for your MySQL server (including RDS, Aurora, and Aurora Serverless) to use `serverless-mysql`. You should just be aware of the following two scenarios.
 
@@ -391,5 +400,25 @@ Other tests that use larger configurations were extremely successful too, but I'
 ## Contributions
 Contributions, ideas and bug reports are welcome and greatly appreciated. Please add [issues](https://github.com/jeremydaly/serverless-mysql/issues) for suggestions and bug reports or create a pull request.
 
-## TODO
-- Add connection retries on failed queries
+## Query Retries
+The module supports automatic retries for transient query errors. When a query fails with a retryable error (such as deadlocks, timeouts, or connection issues), the module will automatically retry the query using the configured backoff strategy.
+
+By default, query retries are disabled (maxQueryRetries = 0) for backward compatibility with previous versions. To enable this feature, set `maxQueryRetries` to a value greater than 0.
+
+You can configure the maximum number of retries with the `maxQueryRetries` option (default: 0) and the backoff strategy with the `queryRetryBackoff` option (default: 'full'). The module will use the same backoff algorithms as for connection retries.
+
+```javascript
+const mysql = require('serverless-mysql')({
+  config: {
+    host: process.env.ENDPOINT,
+    database: process.env.DATABASE,
+    user: process.env.USERNAME,
+    password: process.env.PASSWORD
+  },
+  maxQueryRetries: 5, // Enable query retries with 5 maximum attempts
+  queryRetryBackoff: 'decorrelated',
+  onQueryRetry: (err, retries, delay, type) => {
+    console.log(`Retrying query after error: ${err.code}, attempt: ${retries}, delay: ${delay}ms`)
+  }
+})
+```

index.d.ts

@@ -92,6 +92,18 @@ declare namespace serverlessMysql {
      * This also attaches the SQL query to error objects when a query fails, making it easier to debug.  false
      */
     returnFinalSqlQuery?: boolean;
+    /**
+     * Integer  Maximum number of times to retry a query before giving up.  0
+     */
+    maxQueryRetries?: number;
+    /**
+     * String or Function  Backoff algorithm to be used when retrying queries. Possible values are full and decorrelated, or you can also specify your own algorithm.  full
+     */
+    queryRetryBackoff?: string | Function;
+    /**
+     * function  Event callback when queries are retried.
+     */
+    onQueryRetry?: Function;
   };
 
   class Transaction {

index.js

@@ -33,10 +33,27 @@ module.exports = (params) => {
     'ETIMEDOUT' // if the connection times out
   ]
 
+  // Common Transient Query Errors that can be retried
+  const retryableQueryErrors = [
+    'ER_LOCK_DEADLOCK', // Deadlock found when trying to get lock
+    'ER_LOCK_WAIT_TIMEOUT', // Lock wait timeout exceeded
+    'ER_QUERY_INTERRUPTED', // Query execution was interrupted
+    'ER_QUERY_TIMEOUT', // Query execution time exceeded
+    'ER_CONNECTION_KILLED', // Connection was killed
+    'ER_LOCKING_SERVICE_TIMEOUT', // Locking service timeout
+    'ER_LOCKING_SERVICE_DEADLOCK', // Locking service deadlock
+    'ER_ABORTING_CONNECTION', // Aborted connection
+    'PROTOCOL_CONNECTION_LOST', // Connection lost
+    'PROTOCOL_SEQUENCE_TIMEOUT', // Connection timeout
+    'ETIMEDOUT', // Connection timeout
+    'ECONNRESET' // Connection reset
+  ]
+
   // Init setting values
   let MYSQL, manageConns, cap, base, maxRetries, connUtilization, backoff,
     zombieMinTimeout, zombieMaxTimeout, maxConnsFreq, usedConnsFreq,
-    onConnect, onConnectError, onRetry, onClose, onError, onKill, onKillError, PromiseLibrary, returnFinalSqlQuery
+    onConnect, onConnectError, onRetry, onClose, onError, onKill, onKillError, PromiseLibrary, returnFinalSqlQuery,
+    maxQueryRetries, onQueryRetry, queryRetryBackoff
 
   /********************************************************************/
   /**  HELPER/CONVENIENCE FUNCTIONS                                  **/
@@ -209,57 +226,79 @@ module.exports = (params) => {
 
   // Main query function
   const query = async function (...args) {
-
     // Establish connection
     await connect()
 
-    // Run the query
-    return new PromiseLibrary((resolve, reject) => {
-      if (client !== null) {
-        // If no args are passed in a transaction, ignore query
-        if (this && this.rollback && args.length === 0) { return resolve([]) }
+    // Track query retries
+    let queryRetries = 0
 
-        const queryObj = client.query(...args, async (err, results) => {
-          if (returnFinalSqlQuery && queryObj.sql && err) {
-            err.sql = queryObj.sql
-          }
+    // Function to execute the query with retry logic
+    const executeQuery = async () => {
+      return new PromiseLibrary((resolve, reject) => {
+        if (client !== null) {
+          // If no args are passed in a transaction, ignore query
+          if (this && this.rollback && args.length === 0) { return resolve([]) }
 
-          if (err && err.code === 'PROTOCOL_SEQUENCE_TIMEOUT') {
-            client.destroy() // destroy connection on timeout
-            resetClient() // reset the client
-            reject(err) // reject the promise with the error
-          } else if (
-            err && (/^PROTOCOL_ENQUEUE_AFTER_/.test(err.code)
-              || err.code === 'PROTOCOL_CONNECTION_LOST'
-              || err.code === 'EPIPE'
-              || err.code === 'ECONNRESET')
-          ) {
-            resetClient() // reset the client
-            return resolve(query(...args)) // attempt the query again
-          } else if (err) {
-            if (this && this.rollback) {
-              await query('ROLLBACK')
-              this.rollback(err)
+          const queryObj = client.query(...args, async (err, results) => {
+            if (returnFinalSqlQuery && queryObj.sql && err) {
+              err.sql = queryObj.sql
             }
-            reject(err)
-          }
 
-          if (returnFinalSqlQuery && queryObj.sql) {
-            if (Array.isArray(results)) {
-              Object.defineProperty(results, 'sql', {
-                enumerable: false,
-                value: queryObj.sql
-              })
-            } else if (results && typeof results === 'object') {
-              results.sql = queryObj.sql
+            if (err && err.code === 'PROTOCOL_SEQUENCE_TIMEOUT') {
+              client.destroy() // destroy connection on timeout
+              resetClient() // reset the client
+              reject(err) // reject the promise with the error
+            } else if (
+              err && (/^PROTOCOL_ENQUEUE_AFTER_/.test(err.code)
+                || err.code === 'PROTOCOL_CONNECTION_LOST'
+                || err.code === 'EPIPE'
+                || err.code === 'ECONNRESET')
+            ) {
+              resetClient() // reset the client
+              return resolve(query(...args)) // attempt the query again
+            } else if (err && retryableQueryErrors.includes(err.code) && queryRetries < maxQueryRetries) {
+              // Increment retry counter
+              queryRetries++
+
+              // Calculate backoff time
+              let wait = 0
+              let sleep = queryRetryBackoff === 'decorrelated' ? decorrelatedJitter(wait) :
+                typeof queryRetryBackoff === 'function' ? queryRetryBackoff(wait, queryRetries) :
+                  fullJitter()
+
+              // Fire onQueryRetry event
+              onQueryRetry(err, queryRetries, sleep, typeof queryRetryBackoff === 'function' ? 'custom' : queryRetryBackoff)
+
+              // Wait and retry
+              await delay(sleep)
+              return resolve(executeQuery())
+            } else if (err) {
+              if (this && this.rollback) {
+                await query('ROLLBACK')
+                this.rollback(err)
+              }
+              reject(err)
             }
-          }
 
-          return resolve(results)
-        })
-      }
-    })
+            if (returnFinalSqlQuery && queryObj.sql) {
+              if (Array.isArray(results)) {
+                Object.defineProperty(results, 'sql', {
+                  enumerable: false,
+                  value: queryObj.sql
+                })
+              } else if (results && typeof results === 'object') {
+                results.sql = queryObj.sql
+              }
+            }
 
+            return resolve(results)
+          })
+        }
+      })
+    }
+
+    // Execute the query with retry logic
+    return executeQuery()
   } // end query
 
   // Change user method
@@ -449,6 +488,12 @@ module.exports = (params) => {
   usedConnsFreq = Number.isInteger(cfg.usedConnsFreq) ? cfg.usedConnsFreq : 0 // default to 0 ms
   returnFinalSqlQuery = cfg.returnFinalSqlQuery === true // default to false
 
+  // Query retry settings
+  maxQueryRetries = Number.isInteger(cfg.maxQueryRetries) ? cfg.maxQueryRetries : 0 // default to 0 attempts (disabled for backward compatibility)
+  queryRetryBackoff = typeof cfg.queryRetryBackoff === 'function' ? cfg.queryRetryBackoff :
+    cfg.queryRetryBackoff && ['full', 'decorrelated'].includes(cfg.queryRetryBackoff.toLowerCase()) ?
+      cfg.queryRetryBackoff.toLowerCase() : 'full' // default to full Jitter
+
   // Event handlers
   onConnect = typeof cfg.onConnect === 'function' ? cfg.onConnect : () => { }
   onConnectError = typeof cfg.onConnectError === 'function' ? cfg.onConnectError : () => { }
@@ -457,6 +502,7 @@ module.exports = (params) => {
   onError = typeof cfg.onError === 'function' ? cfg.onError : () => { }
   onKill = typeof cfg.onKill === 'function' ? cfg.onKill : () => { }
   onKillError = typeof cfg.onKillError === 'function' ? cfg.onKillError : () => { }
+  onQueryRetry = typeof cfg.onQueryRetry === 'function' ? cfg.onQueryRetry : () => { }
 
   let connCfg = {}
 

test/integration/query-retries.spec.js

@@ -0,0 +1,138 @@
+'use strict';
+
+const { expect } = require('chai');
+const sinon = require('sinon');
+const {
+    createTestConnection,
+    setupTestTable,
+    cleanupTestTable,
+    closeConnection
+} = require('./helpers/setup');
+
+describe('Query Retries Integration Tests', function () {
+    this.timeout(15000);
+
+    let db;
+    let originalQuery;
+    let queryStub;
+    let onQueryRetrySpy;
+    const TEST_TABLE = 'retry_test_table';
+    const TABLE_SCHEMA = `
+    id INT AUTO_INCREMENT PRIMARY KEY,
+    name VARCHAR(255) NOT NULL,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+  `;
+
+    beforeEach(async function () {
+        onQueryRetrySpy = sinon.spy();
+
+        db = createTestConnection({
+            maxQueryRetries: 3,
+            onQueryRetry: onQueryRetrySpy
+        });
+
+        await setupTestTable(db, TEST_TABLE, TABLE_SCHEMA);
+
+        // Store the original query method
+        originalQuery = db.getClient().query;
+    });
+
+    afterEach(async function () {
+        // Restore the original query method if it was stubbed
+        if (queryStub && queryStub.restore) {
+            queryStub.restore();
+        }
+
+        try {
+            await cleanupTestTable(db, TEST_TABLE);
+
+            if (db) {
+                await db.end({ timeout: 5000 });
+                await closeConnection(db);
+            }
+        } catch (err) {
+            console.error('Error during cleanup:', err);
+        }
+    });
+
+    it('should retry queries that fail with retryable errors', async function () {
+        // Create a counter to track the number of query attempts
+        let attempts = 0;
+
+        // Stub the client's query method to simulate a deadlock error on first attempt
+        queryStub = sinon.stub(db.getClient(), 'query').callsFake(function (sql, values, callback) {
+            attempts++;
+
+            // If this is the first or second attempt, simulate a deadlock error
+            if (attempts <= 2) {
+                const error = new Error('Deadlock found when trying to get lock');
+                error.code = 'ER_LOCK_DEADLOCK';
+
+                // Call the callback with the error
+                if (typeof values === 'function') {
+                    values(error);
+                } else {
+                    callback(error);
+                }
+
+                // Return a mock query object
+                return { sql: typeof sql === 'string' ? sql : sql.sql };
+            }
+
+            // On the third attempt, succeed
+            return originalQuery.apply(this, arguments);
+        });
+
+        // Execute a query that should be retried
+        await db.query('INSERT INTO retry_test_table (name) VALUES (?)', ['Test Retry']);
+
+        // Verify the query was attempted multiple times
+        expect(attempts).to.be.at.least(3);
+
+        // Verify the onQueryRetry callback was called
+        expect(onQueryRetrySpy.callCount).to.equal(2);
+
+        // Verify the data was actually inserted
+        const result = await db.query('SELECT * FROM retry_test_table WHERE name = ?', ['Test Retry']);
+        expect(result).to.have.lengthOf(1);
+        expect(result[0].name).to.equal('Test Retry');
+    });
+
+    it('should give up after maxQueryRetries attempts', async function () {
+        // Create a counter to track the number of query attempts
+        let attempts = 0;
+
+        // Stub the client's query method to always fail with a retryable error
+        queryStub = sinon.stub(db.getClient(), 'query').callsFake(function (sql, values, callback) {
+            attempts++;
+
+            const error = new Error('Lock wait timeout exceeded');
+            error.code = 'ER_LOCK_WAIT_TIMEOUT';
+
+            // Call the callback with the error
+            if (typeof values === 'function') {
+                values(error);
+            } else {
+                callback(error);
+            }
+
+            // Return a mock query object
+            return { sql: typeof sql === 'string' ? sql : sql.sql };
+        });
+
+        // Execute a query that should be retried but eventually fail
+        try {
+            await db.query('INSERT INTO retry_test_table (name) VALUES (?)', ['Should Fail']);
+            expect.fail('Query should have failed after max retries');
+        } catch (error) {
+            // Verify the error is the expected one
+            expect(error.code).to.equal('ER_LOCK_WAIT_TIMEOUT');
+
+            // Verify the query was attempted the maximum number of times (initial + retries)
+            expect(attempts).to.equal(4); // 1 initial + 3 retries
+
+            // Verify the onQueryRetry callback was called the expected number of times
+            expect(onQueryRetrySpy.callCount).to.equal(3);
+        }
+    });
+});