Docserver: Make ObjectStore set calls return Futures

chrome/common/extensions/docs/server2/persistent_object_store.py

Index: chrome/common/extensions/docs/server2/persistent_object_store.py
diff --git a/chrome/common/extensions/docs/server2/persistent_object_store.py b/chrome/common/extensions/docs/server2/persistent_object_store.py
index 0feb6d81beb56e9a926433ebc30ae48f0a54142c..c0ae8734c01cdaeab51f7c1ada93c1b8b2f35dd5 100644
--- a/chrome/common/extensions/docs/server2/persistent_object_store.py
+++ b/chrome/common/extensions/docs/server2/persistent_object_store.py
@@ -5,7 +5,7 @@
 from appengine_wrappers import db
 from datastore_models import PersistentObjectStoreItem
 from environment import IsDevServer
-from future import Future
+from future import All, Future
 from object_store import ObjectStore
 
 
@@ -16,14 +16,14 @@ class PersistentObjectStore(ObjectStore):
     self._namespace = namespace
 
   def SetMulti(self, mapping):
-    futures = []
-    for key, value in mapping.items():
-      futures.append(db.put_async(
-          PersistentObjectStoreItem.CreateItem(self._namespace, key, value)))
+    rpcs = [db.put_async(
+        PersistentObjectStoreItem.CreateItem(self._namespace, key, value))
+        for key, value in mapping.iteritems()]
     # If running the dev server, the futures don't complete until the server is
     # *quitting*. This is annoying. Flush now.
     if IsDevServer():
-      [future.wait() for future in futures]
+      [rpc.wait() for rpc in rpcs]
+    return All(Future(callback=lambda: rpc.get_result()) for rpc in rpcs)
 
   def GetMulti(self, keys):
     db_futures = dict(