[Extensions DocServer]: More datastore API updates

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

 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import cPickle
-import googledatastore as datastore
 import logging
 
 from future import Future
+from gcloud import datastore
 
 # N.B.: In order to use this module you should have a working cloud development
 # environment configured with the googledatastore module installed.
 #
 # Please see https://cloud.google.com/datastore/docs/getstarted/start_python/
 
 
-_DATASET_NAME = 'chrome-apps-doc'
+_PROJECT_NAME = 'chrome-apps-doc'
 _PERSISTENT_OBJECT_KIND = 'PersistentObjectStoreItem'
 _VALUE_PROPERTY_NAME = 'pickled_value'
 
 # The max number of entities to include in a single request. This is capped at
 # 500 by the service. In practice we may send fewer due to _MAX_REQUEST_SIZE
 _MAX_BATCH_SIZE = 500
 
 
 # The maximum entity size allowed by Datastore.
 _MAX_ENTITY_SIZE = 1024*1024
 
 
 # The maximum request size (in bytes) to send Datastore. This is an approximate
 # size based on the sum of entity blob_value sizes.
 _MAX_REQUEST_SIZE = 5*1024*1024
 
 
-def _CreateEntity(name, value):
-  entity = datastore.Entity(exclude_from_indexes=[_VALUE_PROPERTY_NAME])
-  path = entity.key.path.add()
-  path.kind = _PERSISTENT_OBJECT_KIND
-  path.name = name
-  entity.update({_VALUE_PROPERTY_NAME: value})
+def _CreateEntity(client, name, value):
+  key = client.key(_PERSISTENT_OBJECT_KIND, name)
+  entity = datastore.Entity(
+      key=key, exclude_from_indexes=[_VALUE_PROPERTY_NAME])
+  entity[_VALUE_PROPERTY_NAME] = value
   return entity
 
 
-def _CreateBatches(data):
+def _CreateBatches(client, data):
   '''Constructs batches of at most _MAX_BATCH_SIZE entities to cover all
   entities defined in |data| without exceeding the transaction size limit.
   This is a generator emitting lists of entities.
   '''
   def get_size(entity):
-    return len(entity.properties[_VALUE_PROPERTY_NAME].value.blob_value)
+    return len(entity[_VALUE_PROPERTY_NAME])
 
-  entities = [_CreateEntity(name, value) for name, value in data.iteritems()]
+  entities = [_CreateEntity(client, name, value)
+              for name, value in data.iteritems()]
   batch_start = 0
   batch_end = 1
   batch_size = get_size(entities[0])
   while batch_end < len(entities):
     next_size = get_size(entities[batch_end])
     if (batch_size + next_size > _MAX_REQUEST_SIZE or
         batch_end - batch_start >= _MAX_BATCH_SIZE):
       yield entities[batch_start:batch_end], batch_end, len(entities)
       batch_start = batch_end
       batch_size = 0
@@ skipping 34 matching lines @@
     Key('PersistentObjectStoreItem', 'animal/dog')
     Key('PersistentObjectStoreItem', 'animal/trolling cat')
 
   If given |original_data|, this will only push key-value pairs for entries that
   are either new or have changed from their original (pickled) value.
 
   Caveat: Pickling and unpickling a dictionary can (but does not always) change
   its key order. This means that objects will often be seen as changed even when
   they haven't changed.
   '''
-  datastore.set_options(dataset=_DATASET_NAME)
+  client = datastore.Client(_PROJECT_NAME)
 
   def flatten(dataset):
     flat = {}
     for namespace, items in dataset.iteritems():
       for k, v in items.iteritems():
         flat['%s/%s' % (namespace, k)] = cPickle.dumps(v)
     return flat
 
   logging.info('Flattening data sets...')
   data = flatten(data)
   original_data = flatten(original_data)
 
   logging.info('Culling new data...')
   for k in data.keys():
     if ((k in original_data and original_data[k] == data[k]) or
         (len(data[k]) > _MAX_ENTITY_SIZE)):
       del data[k]
 
-  for batch, n, total in _CreateBatches(data):
-    commit_request = datastore.CommitRequest()
-    commit_request.mode = datastore.CommitRequest.NON_TRANSACTIONAL
-    commit_request.mutations.upsert.extend(list(batch))
-
+  for entities, n, total in _CreateBatches(client, data):
+    batch = client.batch()
+    for e in entities:
+      batch.put(e)
     logging.info('Committing %s/%s entities...' % (n, total))
-    datastore.commit(commit_request)
+    batch.commit()