datastore: support deferred results. fixes googleapis#77.

stephenplusplus · Aug 23, 2014 · 53db768 · 53db768
1 parent 00553cc
commit 53db768
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 5 deletions.
diff --git a/lib/datastore/entity.js b/lib/datastore/entity.js
@@ -95,7 +95,10 @@ var SIGN_TO_ORDER = {
  *
  * @example
  * ```js
- * var key = new Key('Company', 123);
+ * var key = new Key({
+ *   namespace: 'ns',
+ *   path: ['Company', 123]
+ * });
  * ```
  */
 function Key(options) {

diff --git a/lib/datastore/transaction.js b/lib/datastore/transaction.js
@@ -222,11 +222,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
@@ -78,6 +78,41 @@ describe('Dataset', function() {
     });
   });
 
+  it('should continue looking for deferred results', function(done) {
+    var ds = new datastore.Dataset({ projectId: 'test' });
+    ds.transaction.makeReq = function(method, proto, typ, callback) {
+      assert.equal(method, 'lookup');
+      if (mockRespGet.deferred.length) {
+        // Revert deferred to original state.
+        mockRespGet.deferred = [];
+      } else {
+        mockRespGet.deferred = [
+          {
+            partition_id: {
+              dataset_id: 's~bamboo-shift-504',
+              namespace: null
+            },
+            path_element: [
+              {
+                kind: 'Kind',
+                id: 5732568548769792,
+                name: null
+              }
+            ]
+          }
+        ];
+      }
+      callback(null, mockRespGet);
+    };
+    var key = ds.key('Kind', 5732568548769792);
+    ds.get([key], function(err, entities) {
+      assert.equal(entities.length, 2);
+      assert.equal(entities[0].key.path_[1], 5732568548769792);
+      assert.equal(entities[1].key.path_[1], 5732568548769792);
+      done();
+    });
+  });
+
   it('should delete by key', function(done) {
     var ds = new datastore.Dataset({ projectId: 'test' });
     ds.transaction.makeReq = function(method, proto, typ, callback) {