Add backpressure functionality to row stream (#505)

* Add backpressure functionality to row stream

* Fix indexing of row buffer

* Increment rowIndex before if statement

* Add tests for highWaterMark

* Rename test suite name to follow naming convention

* Add try/catch and pass error to callback

* Refactor tests into a loop

* Add check to ensure rowIndex is greater than 0 before decrement

Author: sfc-gh-ext-simba-lf

lib/connection/result/row_stream.js

@@ -173,7 +173,37 @@ function RowStream(statement, context, options)
     // add the next row to the read queue
     process.nextTick(function ()
     {
-      self.push(externalizeRow(row, columns, mapColumnIdToExtractFnName));
+      // check if there are still rows available in the rowBuffer
+      if (rowBuffer && rowIndex > 0)
+      {
+        rowIndex--; // decrement the index to include the previous row in the while loop
+
+        // push() data to readable stream until highWaterMark threshold is reached or all rows are pushed
+        while (rowIndex < rowBuffer.length)
+        {
+          row = rowBuffer[rowIndex++];
+
+          // if buffer has reached the threshold based on the highWaterMark value then
+          // push() will return false and pause sending data to the buffer until the data is read from the buffer
+          if (!self.push(externalizeRow(row, columns, mapColumnIdToExtractFnName)))
+          {
+            break;
+          }
+        }
+
+        // check if all rows in rowBuffer has been pushed to the readable stream
+        if (rowIndex === rowBuffer.length)
+        {
+          // reset the buffer and index
+          rowBuffer = null;
+          rowIndex = 0;
+        }
+      }
+      else // No more rows left in the buffer
+      {
+        // Push the last row in the buffer
+        self.push(externalizeRow(row, columns, mapColumnIdToExtractFnName));
+      }
     });
   };
 

lib/connection/statement.js

@@ -678,9 +678,15 @@ function invokeStatementComplete(statement, context)
       // array to the complete callback as the last argument
       var rows = [];
       statement.streamRows()
-        .on('data', function (row)
+        .on('readable', function () // read only when data is available
         {
-          rows.push(row);
+          let row;
+
+          // while there are rows available to read, push row to results array
+          while ((row = this.read()) !== null)
+          {
+            rows.push(row);
+          }
         })
         .on('end', function ()
         {

test/integration/testStreamRows.js

@@ -2,6 +2,7 @@
  * Copyright (c) 2015-2019 Snowflake Computing Inc. All rights reserved.
  */
 var assert = require('assert');
+var async = require('async');
 var testUtil = require('./testUtil');
 require('events').EventEmitter.prototype._maxListeners = 100;
 
@@ -372,4 +373,92 @@ describe('Test Stream Rows API', function ()
       }
     });
   });*/
+
+});
+
+describe('Test Stream Rows HighWaterMark', function ()
+{
+  this.timeout(300000);
+
+  before(function (done)
+  {
+    connection = testUtil.createConnection();
+    testUtil.connect(connection, done);
+  });
+
+  after(function (done)
+  {
+    testUtil.destroyConnection(connection, done);
+  });
+
+  var testingFunc = function (highWaterMark, expectedRowCount, callback)
+  {
+    async.series(
+      [
+        function (callback)
+        {
+          // select table with row count equal to expectedRowCount
+          var statement = connection.execute({
+            sqlText: `SELECT seq8() FROM table(generator(rowCount => ${expectedRowCount}));`,
+            streamResult: true,
+            complete: function ()
+            {
+              var actualRowCount = 0;
+              var rowIndex;
+
+              var stream = statement.streamRows();
+              stream.on('error', function (err)
+              {
+                callback(err);
+              });
+              stream.on('readable', function ()
+              {
+                rowIndex = 0;
+
+                while (this.read() !== null)
+                {
+                  actualRowCount++;
+                  rowIndex++;
+                }
+
+                // assert the amount of rows read per loop never exceeds the highWaterMark threshold
+                try
+                {
+                  assert.ok(rowIndex <= highWaterMark);
+                }
+                catch (err)
+                {
+                  stream.destroy(err); // passes error to the stream error event
+                }
+              });
+              stream.on('end', function ()
+              {
+                try
+                {
+                  // assert the total number of rows is equal to the specified row count
+                  assert.strictEqual(actualRowCount, expectedRowCount);
+                  callback();
+                }
+                catch (err)
+                {
+                  callback(err);
+                }
+              });
+            }
+          });
+        }
+      ],
+      callback
+    );
+  };
+
+  const highWaterMarkValue = 10; // default parameter value is 10 (based on PARAM_ROW_STREAM_HIGH_WATER_MARK)
+
+  [1000, 10000, 100000, 1000000].forEach(rowCount =>
+  {
+    it(`test ${rowCount} rows`, done =>
+    {
+      testingFunc(highWaterMarkValue, rowCount, done);
+    });
+  });
 });