fix stack traces of query() to include the async context (brianc#1762) (brianc#2983)

* fix stack traces of query() to include the async context (brianc#1762)

* rename tests so they are actually run

* conditionally only run async stack trace tests on node 16+

* add stack trace to pg-native

---------

Co-authored-by: Charmander <~@charmander.me>

Author: phiresky

packages/pg-pool/index.js

@@ -39,6 +39,11 @@ function promisify(Promise, callback) {
   const result = new Promise(function (resolve, reject) {
     res = resolve
     rej = reject
+  }).catch(err => {
+    // replace the stack trace that leads to `TCP.onStreamRead` with one that leads back to the
+    // application that created the query
+    Error.captureStackTrace(err);
+    throw err;
   })
   return { callback: cb, result: result }
 }

packages/pg/lib/client.js

@@ -520,6 +520,11 @@ class Client extends EventEmitter {
       if (!query.callback) {
         result = new this._Promise((resolve, reject) => {
           query.callback = (err, res) => (err ? reject(err) : resolve(res))
+        }).catch(err => {
+          // replace the stack trace that leads to `TCP.onStreamRead` with one that leads back to the
+          // application that created the query
+          Error.captureStackTrace(err);
+          throw err;
         })
       }
     }

packages/pg/lib/native/client.js

@@ -174,6 +174,9 @@ Client.prototype.query = function (config, values, callback) {
       result = new this._Promise((resolve, reject) => {
         resolveOut = resolve
         rejectOut = reject
+      }).catch(err => {
+        Error.captureStackTrace(err);
+        throw err;
       })
       query.callback = (err, res) => (err ? rejectOut(err) : resolveOut(res))
     }

packages/pg/test/integration/client/async-stack-trace-tests.js

@@ -0,0 +1,51 @@
+'use strict'
+var helper = require('../test-helper')
+var pg = helper.pg
+
+process.on('unhandledRejection', function (e) {
+  console.error(e, e.stack)
+  process.exit(1)
+})
+
+const suite = new helper.Suite()
+
+// these tests will only work for if --async-stack-traces is on, which is the default starting in node 16.
+const NODE_MAJOR_VERSION = +process.versions.node.split('.')[0];
+if (NODE_MAJOR_VERSION >= 16) {
+  suite.testAsync('promise API async stack trace in pool', async function outerFunction() {
+    async function innerFunction() {
+      const pool = new pg.Pool()
+      await pool.query('SELECT test from nonexistent');
+    }
+    try {
+      await innerFunction();
+      throw Error("should have errored");
+    } catch (e) {
+      const stack = e.stack;
+      if(!e.stack.includes("innerFunction") || !e.stack.includes("outerFunction")) {
+        throw Error("async stack trace does not contain wanted values: " + stack);
+      }
+    }
+  })
+
+  suite.testAsync('promise API async stack trace in client', async function outerFunction() {
+    async function innerFunction() {
+      const client = new pg.Client()
+      await client.connect();
+      try {
+        await client.query('SELECT test from nonexistent');
+      } finally {
+        client.end();
+      }
+    }
+    try {
+      await innerFunction();
+      throw Error("should have errored");
+    } catch (e) {
+      const stack = e.stack;
+      if(!e.stack.includes("innerFunction") || !e.stack.includes("outerFunction")) {
+        throw Error("async stack trace does not contain wanted values: " + stack);
+      }
+    }
+  })
+}