scripts/export_to_gcloud.py

# -*- coding: utf-8 -*-
# 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 dslib


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',
                      default=None,
                      help=('Optional project_id of datastore to write to. If '
                            'not supplied, will be taken from credentials '
                            'file.'))
  parser.add_argument('--namespace', '-n', type=str, action='store',
                      default=None,
                      help='Optional namespace in which to store entities.')
  parser.add_argument('--parent_key', '-p', type=str, action='store',
                      default=None,
                      help='Key of parent entity to insert into. This should '
                      'be in python tuple-literal form, e.g. ("Foo", 1)')
  return parser


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
              entities.
  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)]
    else:
      parent_key = outer_parent_key

    key = datastore.Key(
        kind, idx, project=project_id, parent=parent_key, namespace=namespace)
    e = datastore.Entity(key=key)
    e.update(item)
    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 = dslib.GetClient(creds_file, project_id, namespace)

  if options.parent_key:
    upper_parent_key = c.key(*ast.literal_eval(options.parent_key))
  else:
    upper_parent_key = None


  with open(entities_path, 'r') as f:
    entities = GetEntities(project_id, f, upper_parent_key, namespace)
    dslib.ChunkedBatchWrite(entities, c)