From d59cd15ed28f32271760e4ddf9d9018796fb8b8c Mon Sep 17 00:00:00 2001
From: phiresky <phireskyde+gh@gmail.com>
Date: Wed, 31 May 2023 18:28:53 +0200
Subject: [PATCH] fix stack traces of query() to include the async context
 (#1762) (#2983)

* fix stack traces of query() to include the async context (#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>
---
 packages/pg-pool/index.js                     |  5 ++
 packages/pg/lib/client.js                     |  5 ++
 packages/pg/lib/native/client.js              |  3 ++
 .../client/async-stack-trace-tests.js         | 51 +++++++++++++++++++
 4 files changed, 64 insertions(+)
 create mode 100644 packages/pg/test/integration/client/async-stack-trace-tests.js

diff --git a/packages/pg-pool/index.js b/packages/pg-pool/index.js
index 910aee6d..3c0bf066 100644
--- a/packages/pg-pool/index.js
+++ b/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 }
 }
diff --git a/packages/pg/lib/client.js b/packages/pg/lib/client.js
index 88f2f5f3..a2b7c268 100644
--- a/packages/pg/lib/client.js
+++ b/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;
         })
       }
     }
diff --git a/packages/pg/lib/native/client.js b/packages/pg/lib/native/client.js
index 391a671c..6fb2e060 100644
--- a/packages/pg/lib/native/client.js
+++ b/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))
     }
diff --git a/packages/pg/test/integration/client/async-stack-trace-tests.js b/packages/pg/test/integration/client/async-stack-trace-tests.js
new file mode 100644
index 00000000..d69f7b71
--- /dev/null
+++ b/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);
+      }
+    }
+  })
+}
\ No newline at end of file