From 473f55d874a8975fe5f5aa96105f23a24ef9ffcc Mon Sep 17 00:00:00 2001
From: Stephen Sawchuk <sawchuk@gmail.com>
Date: Wed, 27 Aug 2014 14:30:34 -0400
Subject: [PATCH] datastore: support deferred results. fixes #77.

---
 lib/datastore/entity.js      |  5 ++++-
 lib/datastore/transaction.js | 24 ++++++++++++++++++++----
 test/datastore/dataset.js    | 24 ++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/lib/datastore/entity.js b/lib/datastore/entity.js
index 844c31f7859..5fe63dd2ad6 100644
--- a/lib/datastore/entity.js
+++ b/lib/datastore/entity.js
@@ -70,7 +70,10 @@ var SIGN_TO_ORDER = {
  * @param {string=} options.namespace - Optional namespace.
  *
  * @example
- * var key = new Key('Company', 123);
+ * var key = new Key({
+ *   namespace: 'ns',
+ *   path: ['Company', 123]
+ * });
  */
 function Key(options) {
   this.namespace_ = options.namespace;
diff --git a/lib/datastore/transaction.js b/lib/datastore/transaction.js
index c5896d43340..b6b8bf5c2dc 100644
--- a/lib/datastore/transaction.js
+++ b/lib/datastore/transaction.js
@@ -211,11 +211,27 @@ Transaction.prototype.get = function(keys, callback) {
   }
   this.makeReq('lookup', req, res, function(err, resp) {
     if (err) {
-      return callback(err);
+      callback(err);
+      return;
     }
-    var response = entity.formatArray(resp.found);
-    callback(null, isMultipleRequest ? response : response[0]);
-  });
+    var found = entity.formatArray(resp.found);
+    if (isMultipleRequest && resp.deferred && resp.deferred.length) {
+      // There may be more results. Call `.get` again, and append the results.
+      this.get(
+          resp.deferred.map(entity.keyFromKeyProto), function (err, entities) {
+        if (err) {
+          callback(err);
+          return;
+        }
+        if (resp) {
+          found = (found || []).concat(entities);
+        }
+        callback(null, found);
+      });
+      return;
+    }
+    callback(null, isMultipleRequest ? found : found[0]);
+  }.bind(this));
 };
 
 /**
diff --git a/test/datastore/dataset.js b/test/datastore/dataset.js
index 3af056d9a05..ba441d6dada 100644
--- a/test/datastore/dataset.js
+++ b/test/datastore/dataset.js
@@ -21,6 +21,7 @@
 var assert = require('assert');
 var ByteBuffer = require('bytebuffer');
 var datastore = require('../../lib').datastore;
+var entity = require('../../lib/datastore/entity.js');
 var mockRespGet = require('../testdata/response_get.json');
 var Transaction = require('../../lib/datastore/transaction.js');
 
@@ -78,6 +79,29 @@ describe('Dataset', function() {
     });
   });
 
+  it('should continue looking for deferred results', function(done) {
+    var ds = new datastore.Dataset({ projectId: 'test' });
+    var key = ds.key('Kind', 5732568548769792);
+    var key2 = ds.key('Kind', 5732568548769792);
+    var lookupCount = 0;
+    ds.transaction.makeReq = function(method, proto, typ, callback) {
+      lookupCount++;
+      assert.equal(method, 'lookup');
+      if (mockRespGet.deferred.length) {
+        // Revert deferred to original state.
+        mockRespGet.deferred = [];
+      } else {
+        mockRespGet.deferred = [ entity.keyToKeyProto(key2) ];
+      }
+      callback(null, mockRespGet);
+    };
+    ds.get([key, key2], function(err, entities) {
+      assert.equal(entities.length, 2);
+      assert.equal(lookupCount, 2);
+      done();
+    });
+  });
+
   it('should delete by key', function(done) {
     var ds = new datastore.Dataset({ projectId: 'test' });
     ds.transaction.makeReq = function(method, proto, typ, callback) {