Docserver overhaul: Gitiles away from me.

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

Index: chrome/common/extensions/docs/server2/datastore_util.py
diff --git a/chrome/common/extensions/docs/server2/datastore_util.py b/chrome/common/extensions/docs/server2/datastore_util.py
new file mode 100644
index 0000000000000000000000000000000000000000..2deeb938a26c5e414710f0f370c7c7232d7c042e
--- /dev/null
+++ b/chrome/common/extensions/docs/server2/datastore_util.py
@@ -0,0 +1,135 @@
+# 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.

[Ken Rockot(use gerrit already), 2015/05/26 00:26:24]

This provides the PushData function used exclusively by update_cache.py. In a
nutshell, it pushes a fake datastore into a real datastore. In practice, only
usable from a CE instance or if you locally set DOCSERVER_ACCESS_TOKEN to a
chrome-apps-doc service account token (but nobody should do that).

+
+import cPickle
+import googledatastore as datastore
+import logging
+
+from future import Future
+
+# 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'
+_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

[Ken Rockot(use gerrit already), 2015/05/26 00:26:24]

Couldn't find any documentation for the max request size. 5 MB works
consistently.

+
+
+def _CreateEntity(name, value):
+  entity = datastore.Entity()
+  path = entity.key.path_element.add()
+  path.kind = _PERSISTENT_OBJECT_KIND
+  path.name = name
+  pickled_value_property = entity.property.add()
+  pickled_value_property.name = _VALUE_PROPERTY_NAME
+  pickled_value_property.value.indexed = False
+  pickled_value_property.value.blob_value = value
+  return entity
+
+
+def _CreateBatches(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.property[0].value.blob_value)
+
+  entities = [_CreateEntity(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
+    else:
+      batch_size += next_size
+    batch_end = batch_end + 1
+  if batch_end > batch_start and batch_start < len(entities):
+    yield entities[batch_start:batch_end], batch_end, len(entities)
+
+
+def PushData(data, original_data={}):
+  '''Pushes a bunch of data into the datastore. The data should be a dict. Each
+  key is treated as a namespace, and each value is also a dict. A new datastore
+  entry is upserted for every inner key, with the value pickled into the
+  |pickled_value| field.
+
+  For example, if given the dictionary:
+
+  {
+    'fruit': {
+      'apple': 1234,
+      'banana': 'yellow',
+      'trolling carrot': { 'arbitrarily complex': ['value', 'goes', 'here'] }
+    },
+    'animal': {
+      'sheep': 'baaah',
+      'dog': 'woof',
+      'trolling cat': 'moo'
+    }
+  }
+
+  this would result in a push of 6 keys in total, with the following IDs:
+
+    Key('PersistentObjectStoreItem', 'fruit/apple')
+    Key('PersistentObjectStoreItem', 'fruit/banana')
+    Key('PersistentObjectStoreItem', 'fruit/trolling carrot')
+    Key('PersistentObjectStoreItem', 'animal/sheep')
+    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)
+
+  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]

[Ken Rockot(use gerrit already), 2015/05/26 00:26:24]

This should be super awesome and get us very tiny incremental updates (at which
point we should also do selective memcache eviction rather than flush_all).
However, one small problem: dicts do not retain key ordering across pickles. :(

So this is not always true:

  cPickle.dumps(cPickle.loads(cPickle.dumps(mydict))) == cPickle.dumps(mydict)

Because of this, every update still pushes about 1200 keys into the datastore.
We should fix it, but it doesn't seem like a major issue.

[not at google - send to devlin, 2015/06/04 22:40:45]

What about an OrderedDict? That should pickle, and it should unpickle in the
same order. It's a PITA to need to use OrderedDict everywhere, but if it's super
awesome to get incremental updates, worth a shot.

or you can maintain a list on the side.

I haven't dived deep into this code, so I might be totally off here.

[Ken Rockot(use gerrit already), 2015/06/05 00:21:50]

No, I think that's the right approach. I didn't do OrderedDict in this CL
because it wasn't strictly necessary and it would make the CL even bigger. Nice
thing is, we can always move bits over to OrderedDict incrementally. I'll try it
with maybe CompiledFileSystem soon.

+
+  for batch, n, total in _CreateBatches(data):
+    commit_request = datastore.CommitRequest()
+    commit_request.mode = datastore.CommitRequest.NON_TRANSACTIONAL
+    commit_request.mutation.upsert.extend(list(batch))
+
+    logging.info('Committing %s/%s entities...' % (n, total))
+    datastore.commit(commit_request)