spanner: Properly handle multi-row mutations with heterogeneous columns

packages/spanner/package.json

@@ -35,6 +35,7 @@
   "dependencies": {
     "@google-cloud/common": "^0.13.0",
     "@google-cloud/common-grpc": "^0.3.0",
+    "array-uniq": "^1.0.3",
     "arrify": "^1.0.1",
     "checkpoint-stream": "^0.1.0",
     "events-intercept": "^2.0.0",
@@ -44,6 +45,7 @@
     "google-proto-files": "^0.12.0",
     "is": "^3.1.0",
     "lodash.chunk": "^4.2.0",
+    "lodash.flatten": "^4.4.0",
     "lodash.snakecase": "^4.1.1",
     "merge-stream": "^1.0.1",
     "methmeth": "^1.1.0",

packages/spanner/src/transaction-request.js

@@ -23,7 +23,9 @@
 var arrify = require('arrify');
 var common = require('@google-cloud/common');
 var extend = require('extend');
+var flatten = require('lodash.flatten');
 var is = require('is');
+var uniq = require('array-uniq');
 
 /**
  * @type {module:spanner/codec}
@@ -682,16 +684,28 @@ TransactionRequest.prototype.upsert = function(table, keyVals, callback) {
 TransactionRequest.prototype.mutate_ = function(method, table, keyVals, cb) {
   keyVals = arrify(keyVals);
 
-  var mutation = {};
+  var columns = uniq(flatten(keyVals.map(Object.keys))).sort();
 
-  mutation[method] = {
-    table: table,
-    columns: Object.keys(keyVals[0]),
-    values: keyVals.map(function(keyVal) {
-      return Object.keys(keyVal).map(function(key) {
-        return codec.encode(keyVal[key]);
-      });
-    })
+  var values = keyVals.map(function(keyVal, index) {
+    var keys = Object.keys(keyVal);
+
+    var missingColumns = columns.filter(column => keys.indexOf(column) === -1);
+
+    if (missingColumns.length > 0) {
+      throw new Error([
+        `Row at index ${index} does not contain the correct number of columns.`,
+        `Missing columns: ${JSON.stringify(missingColumns)}`
+      ].join('\n\n'));
+    }
+
+    return columns.map(function(column) {
+      var value = keyVal[column];
+      return codec.encode(value);
+    });
+  });
+
+  var mutation = {
+    [method]: { table, columns, values }
   };
 
   if (this.transaction) {

packages/spanner/system-test/spanner.js

@@ -108,6 +108,47 @@ var spanner = new Spanner(env);
       }, execAfterOperationComplete(done));
     });
 
+    describe('uneven rows', function() {
+      it('should allow differently-ordered rows', function(done) {
+        var data = [
+          {
+            Key: generateName('id'),
+            BoolValue: true,
+            IntValue: spanner.int(10)
+          },
+          {
+            Key: generateName('id'),
+            IntValue: spanner.int(10),
+            BoolValue: true
+          }
+        ];
+
+        table.insert(data, function(err) {
+          assert.ifError(err);
+
+          database.run({
+            sql: `SELECT * FROM \`${table.name}\` WHERE Key = @a OR KEY = @b`,
+            params: {
+              a: data[0].Key,
+              b: data[1].Key
+            }
+          }, function(err, rows) {
+            assert.ifError(err);
+
+            var row1 = rows[0].toJSON();
+            assert.deepStrictEqual(row1.IntValue, data[0].IntValue);
+            assert.deepStrictEqual(row1.BoolValue, data[0].BoolValue);
+
+            var row2 = rows[1].toJSON();
+            assert.deepStrictEqual(row2.IntValue, data[1].IntValue);
+            assert.deepStrictEqual(row2.BoolValue, data[1].BoolValue);
+
+            done();
+          });
+        });
+      });
+    });
+
     describe('structs', function() {
       it('should correctly decode structs', function(done) {
         var query = 'SELECT ARRAY(SELECT as struct 1, "hello")';

packages/spanner/test/transaction-request.js

@@ -759,26 +759,44 @@ describe('TransactionRequest', function() {
   describe('mutate_', function() {
     var METHOD = 'methodName';
     var TABLE = 'table-name';
-    var KEYVALS = { key: 'value' };
-
-    var ENCODED_VALUE = {
-      encoded: true
-    };
+    var KEYVALS = [
+      {
+        key: '1-key-value',
+        anotherNullable: '1-anotherNullable-value',
+        nonNullable: '1-nonNullable-value',
+        nullable: '1-nullable-value'
+      },
+      { /* keys defined in different order */
+        key: '2-key-value',
+        nullable: null,
+        nonNullable: '2-nonNullable-value',
+        anotherNullable: null
+      }
+    ];
 
     var EXPECTED_MUTATION = {};
     EXPECTED_MUTATION[METHOD] = {
       table: TABLE,
-      columns: Object.keys(KEYVALS),
+      columns: ['anotherNullable', 'key', 'nonNullable', 'nullable'],
       values: [
         [
-          ENCODED_VALUE
+          KEYVALS[0].anotherNullable,
+          KEYVALS[0].key,
+          KEYVALS[0].nonNullable,
+          KEYVALS[0].nullable
+        ],
+        [
+          KEYVALS[1].anotherNullable,
+          KEYVALS[1].key,
+          KEYVALS[1].nonNullable,
+          KEYVALS[1].nullable
         ]
       ]
     };
 
     beforeEach(function() {
-      fakeCodec.encode = function() {
-        return ENCODED_VALUE;
+      fakeCodec.encode = function(value) {
+        return value;
       };
     });
 
@@ -787,9 +805,46 @@ describe('TransactionRequest', function() {
 
       function callback() {}
 
-      fakeCodec.encode = function(key) {
-        assert.strictEqual(key, KEYVALS[Object.keys(KEYVALS)[0]]);
-        return ENCODED_VALUE;
+      var numEncodeRequests = 0;
+      fakeCodec.encode = function(value) {
+        numEncodeRequests++;
+
+        switch (numEncodeRequests) {
+          case 1: {
+            assert.strictEqual(value, KEYVALS[0].anotherNullable);
+            break;
+          }
+          case 2: {
+            assert.strictEqual(value, KEYVALS[0].key);
+            break;
+          }
+          case 3: {
+            assert.strictEqual(value, KEYVALS[0].nonNullable);
+            break;
+          }
+          case 4: {
+            assert.strictEqual(value, KEYVALS[0].nullable);
+            break;
+          }
+          case 5: {
+            assert.strictEqual(value, KEYVALS[1].anotherNullable);
+            break;
+          }
+          case 6: {
+            assert.strictEqual(value, KEYVALS[1].key);
+            break;
+          }
+          case 7: {
+            assert.strictEqual(value, KEYVALS[1].nonNullable);
+            break;
+          }
+          case 8: {
+            assert.strictEqual(value, KEYVALS[1].nullable);
+            break;
+          }
+        }
+
+        return value;
       };
 
       var expectedReqOpts = {
@@ -817,6 +872,31 @@ describe('TransactionRequest', function() {
       assert.strictEqual(returnValue, requestReturnValue);
     });
 
+    it('should throw when rows have incorrect amount of columns', function() {
+      var invalidEntry = { key1: 'val' };
+      var caughtError;
+
+      try {
+        transactionRequest.mutate_(METHOD, TABLE, [
+          invalidEntry,
+          { key1: 'val', key2: 'val' }
+        ], assert.ifError);
+      } catch(e) {
+        caughtError = e;
+      } finally {
+        if (!caughtError) {
+          throw new Error('Expected error was not thrown.');
+        }
+
+        var expectedErrorMessage = [
+          'Row at index 0 does not contain the correct number of columns.',
+          'Missing columns: ["key2"]'
+        ].join('\n\n');
+
+        assert.strictEqual(caughtError.message, expectedErrorMessage);
+      }
+    });
+
     it('should push the request to the queue if a transaction', function(done) {
       transactionRequest.transaction = true;