From af4714845603f70e3c1ef635f6c0750ff1987a9e Mon Sep 17 00:00:00 2001
From: nofarham <116541220+nofarham@users.noreply.github.com>
Date: Mon, 22 Jan 2024 22:37:05 +0200
Subject: [PATCH] fix(stream): premature close when using `for await` (#2389)

* Fix node 18 (based on #711)

* Add tests to query stream

* fix async error

* fix indentation
---
 lib/commands/query.js                      | 2 +-
 test/integration/connection/test-stream.js | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/lib/commands/query.js b/lib/commands/query.js
index 0de52ff05a..1b3725d73f 100644
--- a/lib/commands/query.js
+++ b/lib/commands/query.js
@@ -279,7 +279,7 @@ class Query extends Command {
     });
     this.on('end', () => {
       stream.push(null); // pushing null, indicating EOF
-      stream.emit('close'); // notify readers that query has completed
+      setImmediate(() => stream.emit('close')); // notify readers that query has completed
     });
     this.on('fields', fields => {
       stream.emit('fields', fields); // replicate old emitter
diff --git a/test/integration/connection/test-stream.js b/test/integration/connection/test-stream.js
index 133acc8426..57ac6dd5ba 100644
--- a/test/integration/connection/test-stream.js
+++ b/test/integration/connection/test-stream.js
@@ -7,6 +7,7 @@ const assert = require('assert');
 let rows;
 const rows1 = [];
 const rows2 = [];
+const rows3 = [];
 
 connection.query(
   [
@@ -45,7 +46,7 @@ connection.execute(
     }
   }
 );
-connection.execute('SELECT * FROM announcements', (err, _rows) => {
+connection.execute('SELECT * FROM announcements', async (err, _rows) => {
   rows = _rows;
   const s1 = connection.query('SELECT * FROM announcements').stream();
   s1.on('data', row => {
@@ -60,10 +61,15 @@ connection.execute('SELECT * FROM announcements', (err, _rows) => {
       connection.end();
     });
   });
+  const s3 = connection.query('SELECT * FROM announcements').stream();
+  for await (const row of s3) {
+    rows3.push(row);
+  }
 });
 
 process.on('exit', () => {
   assert.deepEqual(rows.length, 2);
   assert.deepEqual(rows, rows1);
   assert.deepEqual(rows, rows2);
+  assert.deepEqual(rows, rows3);
 });