blob: 0ec0c99a7e5bf8d3ca4c9f8117e8ecd6b6307f7d [file] [log] [blame]
# Copyright 2016 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Export entities to gcloud datastore."""
from __future__ import print_function
from gcloud import datastore
import ast
import json
from chromite.lib import commandline
from chromite.lib import iter_utils
def GetParser():
"""Creates the argparse parser."""
parser = commandline.ArgumentParser(description=__doc__)
parser.add_argument('service_acct_json', type=str, action='store',
help='Path to service account credentials JSON file.')
parser.add_argument('entities', type=str, action='store',
help=('Path to file with entities to export. '
'File should be newline-separated JSON entries.'))
parser.add_argument('--project_id', '-i', type=str, action='store',
help=('Optional project_id of datastore to write to. If '
'not supplied, will be taken from credentials '
parser.add_argument('--namespace', '-n', type=str, action='store',
help='Optional namespace in which to store entities.')
parser.add_argument('--parent_key', '-p', type=str, action='store',
help='Key of parent entity to insert into. This should '
'be in python tuple-literal form, e.g. ("Foo", 1)')
return parser
def _GetClient(creds_file, project_id=None, namespace=None):
"""Get a datastore client instance.
creds_file: Path to JSON creds file.
project_id: Optional project_id of datastore. If not supplied,
will use one from creds file.
namespace: Optional namespace to insert into.
if project_id is None:
with open(creds_file, 'r') as f:
project_id = json.load(f)['project_id']
return datastore.Client.from_service_account_json(
creds_file, project_id, namespace=namespace), project_id
class DuplicateKeyError(ValueError):
"""Raised when two Entities have the same key."""
def GetEntities(project_id, json_lines, outer_parent_key=None, namespace=None):
"""Create gcloud entities from json string entries.
project_id: String gcloud project id that entities are for.
json_lines: File or other line-by-line iterator of json strings to turn into
outer_parent_key: Optional datastore.Key instance to act as the parent_key
of all top level entities.
namespace: Optional string namespace for entities.
entity_keys = {}
for line in json_lines:
item = json.loads(line)
kind, idx = item.pop('id')
parent = item.pop('parent', None)
if (kind, idx) in entity_keys:
raise DuplicateKeyError(
'Duplicate entities with id (%s, %s)' % (kind, idx))
if parent:
parent_key = entity_keys[tuple(parent)]
parent_key = outer_parent_key
key = datastore.Key(
kind, idx, project=project_id, parent=parent_key, namespace=namespace)
e = datastore.Entity(key=key)
entity_keys[(kind, idx)] = key
entity_keys[idx] = key
yield e
def main(argv):
parser = GetParser()
options = parser.parse_args(argv)
entities_path = options.entities
creds_file = options.service_acct_json
project_id = options.project_id
namespace = options.namespace
entities = []
c, project_id = _GetClient(creds_file, project_id, namespace)
if options.parent_key:
upper_parent_key = c.key(*ast.literal_eval(options.parent_key))
upper_parent_key = None
with open(entities_path, 'r') as f:
entities = GetEntities(project_id, f, upper_parent_key, namespace)
for chunk in iter_utils.SplitToChunks(entities, _BATCH_CHUNK_SIZE):
batch = c.batch()
for e in chunk: