Rewrite PartialResultStream (#208)

Fixes #180 and #198 by reworking the way we process streamed rows from PartialResultStream on Spanner. The high-level difference is this

- use RowBuilder for all streams. We had two code paths previously
- do not block on returning until we have a non-chunked row. We can hold the final value of a chunked segment and begin returning other data
- handle partial segments. Occasionally we will get a partial row and shouldn't try to send that to the user.

Author: stephenplusplus

handwritten/spanner/package-lock.json

@@ -77,7 +77,6 @@
     "google-gax": "^0.16.0",
     "google-proto-files": "^0.15.0",
     "is": "^3.1.0",
-    "lodash.chunk": "^4.2.0",
     "lodash.flatten": "^4.4.0",
     "lodash.merge": "^4.6.1",
     "lodash.snakecase": "^4.1.1",

handwritten/spanner/package.json

@@ -176,7 +176,7 @@ function decode(value, field) {
     return decoded;
   }
 
-  return decodeValue_(commonGrpc.Service.decodeValue_(value), field.type);
+  return decodeValue_(value, field.type);
 }
 
 codec.decode = decode;

handwritten/spanner/src/codec.js

@@ -17,7 +17,6 @@
 'use strict';
 
 var checkpointStream = require('checkpoint-stream');
-var chunk = require('lodash.chunk');
 var eventsIntercept = require('events-intercept');
 var exec = require('methmeth');
 var extend = require('extend');
@@ -27,7 +26,6 @@ var split = require('split-array-stream').split;
 var streamEvents = require('stream-events');
 var through = require('through2');
 
-var codec = require('./codec.js');
 var RowBuilder = require('./row-builder.js');
 
 /**
@@ -68,45 +66,26 @@ function partialResultStream(requestFn, options) {
     },
   });
 
-  var rowChunks = [];
-  var metadata;
+  var builder;
 
   var userStream = streamEvents(
     through.obj(function(row, _, next) {
-      var formattedRows = [];
-
-      if (row.metadata) {
-        metadata = row.metadata;
-      }
-
-      if (row.chunkedValue) {
-        rowChunks.push(row);
-        next();
-        return;
-      }
-
       if (is.empty(row.values)) {
         next();
         return;
       }
 
-      if (rowChunks.length > 0) {
-        // Done getting all the chunks. Put them together.
-        var builder = new RowBuilder(metadata, rowChunks.concat(row));
-        formattedRows = formattedRows.concat(builder.toJSON());
-        rowChunks.length = 0;
-      } else {
-        var formattedRow = partialResultStream.formatRow_(metadata, row);
-        var multipleRows = is.array(formattedRow[0]);
-
-        if (multipleRows) {
-          formattedRows = formattedRows.concat(formattedRow);
-        } else {
-          formattedRows.push(formattedRow);
-        }
+      // Use RowBuilder to construct and return complete, formatted rows.
+      if (!builder) {
+        builder = new RowBuilder(row.metadata.rowType.fields);
       }
 
-      rowChunks = [];
+      builder.addRow(row);
+
+      // Build the chunks to rows.
+      builder.build();
+
+      var formattedRows = builder.toJSON(builder.flush());
 
       if (options.json) {
         formattedRows = formattedRows.map(exec('toJSON', options.jsonOptions));
@@ -155,41 +134,4 @@ function partialResultStream(requestFn, options) {
     .pipe(userStream);
 }
 
-/**
- * Format a PartialResultSet response from the API. A row object will be created
- * to map each field name to its decoded value.
- *
- * If multiple rows exist in a single PartialResultSet, an array is returned.
- *
- * @param {object} row - A `PartialResultSet` object.
- */
-partialResultStream.formatRow_ = function(metadata, row) {
-  var fields = metadata.rowType.fields;
-
-  if (row.values.length > fields.length) {
-    // More than one row exists. Return an array of formatted rows.
-    var valueSets = chunk(row.values, fields.length);
-
-    return valueSets.map(function(valueSet) {
-      row.values = valueSet;
-      return partialResultStream.formatRow_(metadata, row);
-    });
-  }
-
-  var formattedRow = row.values.map(function(value, index) {
-    var field = fields[index];
-    return {
-      name: field.name,
-      value: codec.decode(value, field),
-    };
-  });
-
-  Object.defineProperty(formattedRow, 'toJSON', {
-    enumerable: false,
-    value: codec.generateToJSONFromRow(formattedRow),
-  });
-
-  return formattedRow;
-};
-
 module.exports = partialResultStream;

handwritten/spanner/src/partial-result-stream.js

@@ -17,6 +17,7 @@
 'use strict';
 
 var codec = require('./codec.js');
+var commonGrpc = require('@google-cloud/common-grpc');
 var is = require('is');
 
 /*!
@@ -25,11 +26,9 @@ var is = require('is');
  * @private
  * @class
  */
-function RowBuilder(metadata, chunks) {
-  this.metadata = metadata;
-  this.fields = this.metadata.rowType.fields;
-  this.chunks = chunks;
-
+function RowBuilder(fields) {
+  this.fields = fields;
+  this.chunks = [];
   this.rows = [[]];
 
   Object.defineProperty(this, 'currentRow', {
@@ -46,7 +45,7 @@ RowBuilder.getValue = function(obj) {
   var value = obj;
 
   if (obj && obj.kind) {
-    value = obj[obj.kind];
+    value = commonGrpc.Service.decodeValue_(obj);
   }
 
   if (value && value.values) {
@@ -72,11 +71,11 @@ RowBuilder.formatValue = function(field, value) {
   }
 
   if (field.code !== 'STRUCT') {
-    return value;
+    return codec.decode(value, field);
   }
 
   return field.structType.fields.reduce(function(struct, field, index) {
-    struct[field.name] = RowBuilder.formatValue(field.type, value[index]);
+    struct[field.name] = RowBuilder.formatValue(field, value[index]);
     return struct;
   }, {});
 };
@@ -116,14 +115,20 @@ RowBuilder.merge = function(type, head, tail) {
   });
 };
 
+/**
+ * Add a PartialResultSet response object to the pending rows.
+ */
+RowBuilder.prototype.addRow = function(row) {
+  this.chunks = this.chunks.concat(row);
+};
+
 /**
  * Appends element to row.
  */
 RowBuilder.prototype.append = function(value) {
   if (this.currentRow.length === this.fields.length) {
     this.rows.push([]);
   }
-
   this.currentRow.push(value);
 };
 
@@ -132,35 +137,59 @@ RowBuilder.prototype.append = function(value) {
  */
 RowBuilder.prototype.build = function() {
   var self = this;
-  var previousChunk;
 
   this.chunks.forEach(function(chunk) {
-    if (previousChunk && previousChunk.chunkedValue) {
-      var type = self.fields[self.currentRow.length - 1].type;
+    // If we have a chunk to merge, merge the values now.
+    if (self.pendingChunk) {
+      var currentColumn = self.currentRow.length % self.fields.length;
       var merged = RowBuilder.merge(
-        type,
-        self.currentRow.pop(),
+        self.fields[currentColumn].type,
+        self.pendingChunk,
         chunk.values.shift()
       );
+      chunk.values = merged.concat(chunk.values);
+      delete self.pendingChunk;
+    }
 
-      merged.forEach(self.append.bind(self));
+    // If the chunk is chunked, store the last value for merging with the next
+    // chunk to be processed.
+    if (chunk.chunkedValue) {
+      self.pendingChunk = chunk.values.pop();
     }
 
     chunk.values.map(RowBuilder.getValue).forEach(self.append.bind(self));
-
-    previousChunk = chunk;
   });
+
+  // As chunks are now in rows, remove them.
+  this.chunks.length = 0;
 };
 
 /**
- * Transforms values into JSON format.
+ * Flush already complete rows.
  */
-RowBuilder.prototype.toJSON = function() {
-  this.build();
+RowBuilder.prototype.flush = function() {
+  var rowsToReturn = this.rows;
+
+  if (
+    !is.empty(this.rows[0]) &&
+    this.currentRow.length !== this.fields.length
+  ) {
+    // Don't return the partial row. Hold onto it for the next iteration.
+    this.rows = this.rows.splice(-1);
+  } else {
+    this.rows = [[]];
+  }
+
+  return rowsToReturn;
+};
 
+/**
+ * Transforms values into JSON format.
+ */
+RowBuilder.prototype.toJSON = function(rows) {
   var fields = this.fields;
 
-  return this.rows.map(function(values) {
+  return rows.map(function(values) {
     var formattedRow = [];
     var serializedRow = {};
 
@@ -169,7 +198,7 @@ RowBuilder.prototype.toJSON = function() {
 
       var column = {
         name: field.name,
-        value: RowBuilder.formatValue(field.type, value),
+        value: RowBuilder.formatValue(field, value),
       };
 
       formattedRow.push(column);

handwritten/spanner/src/row-builder.js

@@ -213,17 +213,11 @@ describe('codec', function() {
       };
     });
 
-    it('should return the value from the common decoder', function() {
+    it('should return the same value if not a special type', function() {
       var value = {};
-      var defaultDecodedValue = {};
-
-      FakeGrpcService.decodeValue_ = function(value_) {
-        assert.strictEqual(value_, value);
-        return defaultDecodedValue;
-      };
 
       var decoded = codec.decode(value, BYPASS_FIELD);
-      assert.strictEqual(decoded, defaultDecodedValue);
+      assert.strictEqual(decoded, value);
     });
 
     it('should return null values as null', function() {