scripts/export_to_gcloud.py

# Copyright 2016 The ChromiumOS Authors
# 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."""

import ast
import json

from google.cloud import datastore  # pylint: disable=import-error

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", encoding="utf-8") as f:
        entities = GetEntities(project_id, f, upper_parent_key, namespace)
        dslib.ChunkedBatchWrite(entities, c)