feat: automatically purge stale prepared statements from cache and retry query

Author: Invader444

lib/base/connection.js

@@ -28,6 +28,7 @@ const Packets = require('../packets/index.js');
 const Commands = require('../commands/index.js');
 const ConnectionConfig = require('../connection_config.js');
 const CharsetToEncoding = require('../constants/charset_encodings.js');
+const { ER_UNKNOWN_STMT_HANDLER } = require('../constants/errors.js');
 const {
   traceCallback,
   tracePromise,
@@ -814,13 +815,45 @@ class BaseConnection extends EventEmitter {
       this.addCommand(executeCommand);
     };
 
+    // We need to intercept and retry prepareAndExecute if we had a stale prepared statement in the cache that the server already released
+    const key = BaseConnection.statementKey(options);
+    const cacheHasKey = this._statements.has(key);
+
     if (executeCommand.onResult) {
       // Callback mode: traceCallback wraps the callback with tracing lifecycle, or calls through directly when no subscribers are registered
       const origExecCb = executeCommand.onResult;
       traceCallback(
         executeChannel,
         (wrappedCb) => {
-          executeCommand.onResult = wrappedCb;
+          if (cacheHasKey) {
+            executeCommand.onResult = (err, ...rest) => {
+              if (err && err.errno === ER_UNKNOWN_STMT_HANDLER) {
+                const origEmit = executeCommand.emit.bind(executeCommand);
+                executeCommand.emit = (eventName, ...rest) => {
+                  if (eventName === 'end') {
+                    // Intercept the 'end' event that will be emitted after this 'error' event is emitted
+                    executeCommand.emit = origEmit;
+                    return false;
+                  }
+
+                  return origEmit(eventName, ...rest);
+                };
+
+                // Listeners may have been added to this execute command, so we re-use it
+                executeCommand.next = null;
+
+                this._statements.delete(key);
+                executeCommand.onResult = wrappedCb;
+                prepareAndExecute(wrappedCb);
+                return;
+              }
+
+              wrappedCb(err, ...rest);
+            };
+          } else {
+            executeCommand.onResult = wrappedCb;
+          }
+
           prepareAndExecute(wrappedCb);
         },
         0,
@@ -837,7 +870,43 @@ class BaseConnection extends EventEmitter {
         null,
         origExecCb
       );
-    } else if (shouldTrace(executeChannel)) {
+
+      return executeCommand;
+    }
+
+    if (cacheHasKey) {
+      const origEmit = executeCommand.emit.bind(executeCommand);
+      executeCommand.emit = (eventName, firstArg, ...rest) => {
+        if (
+          eventName === 'error' &&
+          firstArg &&
+          firstArg.errno === ER_UNKNOWN_STMT_HANDLER
+        ) {
+          executeCommand.emit = (eventName, ...rest) => {
+            if (eventName === 'end') {
+              // Intercept the 'end' event that will be emitted after this 'error' event is emitted
+              executeCommand.emit = origEmit;
+              return false;
+            }
+
+            return origEmit(eventName, ...rest);
+          };
+
+          // Listeners may have been added to this execute command, so we re-use it
+          executeCommand.next = null;
+
+          this._statements.delete(key);
+          prepareAndExecute((err) => {
+            executeCommand.emit('error', err);
+          });
+          return false;
+        }
+
+        return origEmit(eventName, firstArg, ...rest);
+      };
+    }
+
+    if (shouldTrace(executeChannel)) {
       // Event-emitter mode: tracePromise wraps the async lifecycle
       tracePromise(
         executeChannel,

test/integration/connection/test-execute-cached.test.mts

@@ -41,5 +41,55 @@ await describe('Execute Cached', async () => {
     strict.deepEqual(rows3, [{ test: 126 }]);
   });
 
+  await it('should discard cached prepared statements that no longer exist on the server and retry automatically', async () => {
+    // Remove on the server but leave it in the cache
+    // @ts-expect-error: internal access
+    connection._statements.get(key).close();
+
+    // Remember the id
+    // @ts-expect-error: internal access
+    const { id: oldStatementId } = connection._statements.get(key);
+
+    // @ts-expect-error: internal access
+    strict(connection._statements.size === 1);
+    // @ts-expect-error: internal access
+    strict(connection._statements.get(key).query === q);
+    // @ts-expect-error: internal access
+    strict(connection._statements.get(key).parameters.length === 1);
+    // @ts-expect-error: internal access
+    strict(connection._statements.get(key).id === oldStatementId);
+
+    const rows1 = await new Promise<TestRow[]>((resolve, reject) => {
+      connection.execute<TestRow[]>(q, [123], (err, _rows) =>
+        err ? reject(err) : resolve(_rows)
+      );
+    });
+
+    const rows2 = await new Promise<TestRow[]>((resolve, reject) => {
+      connection.execute<TestRow[]>(q, [124], (err, _rows) =>
+        err ? reject(err) : resolve(_rows)
+      );
+    });
+
+    const rows3 = await new Promise<TestRow[]>((resolve, reject) => {
+      connection.execute<TestRow[]>(q, [125], (err, _rows) =>
+        err ? reject(err) : resolve(_rows)
+      );
+    });
+
+    // @ts-expect-error: internal access
+    strict(connection._statements.size === 1);
+    // @ts-expect-error: internal access
+    strict(connection._statements.get(key).query === q);
+    // @ts-expect-error: internal access
+    strict(connection._statements.get(key).parameters.length === 1);
+    // @ts-expect-error: internal access
+    strict(connection._statements.get(key).id !== oldStatementId);
+
+    strict.deepEqual(rows1, [{ test: 124 }]);
+    strict.deepEqual(rows2, [{ test: 125 }]);
+    strict.deepEqual(rows3, [{ test: 126 }]);
+  });
+
   connection.end();
 });