Merge pull request #5 from kieras/develop

0.0.3 version

.gitignore

@@ -1,6 +1,8 @@
 
 # Created by https://www.gitignore.io/api/node,java,linux,macos,maven,python,eclipse,windows,virtualenv,intellij+all
 
+data/
+
 ### Eclipse ###
 
 .metadata

README.md

@@ -5,12 +5,9 @@ Utilities for Google Cloud Datastore.
 
 # Installation
 
-If you don't use `pipsi`, you're missing out.
-Here are [installation instructions](https://github.com/mitsuhiko/pipsi#readme).
-
 Simply run:
 
-    $ pipsi install .
+    $ pipsi install gcdu
 
 
 # Usage

gcdu/__init__.py

@@ -1 +1 @@
-
+# -*- coding: utf-8 -*-

gcdu/__main__.py

@@ -1,3 +1,5 @@
+# -*- coding: utf-8 -*-
+
 from cli import main
 
 main()

gcdu/__version__.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 #   ____    ____   ____    _   _
 #  / ___|  / ___| |  _ \  | | | |
 # | |  _  | |     | | | | | | | |
@@ -7,4 +8,4 @@
 #
 #   Google Cloud Datastore Utils
 
-__version__ = '0.0.1'
+__version__ = '0.0.3'

gcdu/cli.py

@@ -1,5 +1,9 @@
+# -*- coding: utf-8 -*-
 import click
+
 from __version__ import __version__ as version
+from commands.export import export
+from commands.import_cmd import import_cmd
 
 
 @click.group()
@@ -8,66 +12,5 @@ def main():
     """Utilities for Google Cloud Datastore."""
     pass
 
-
-@main.command()
-@click.option('--from-project', '-fp',
-              help='Origin GCP Project.',
-              required=True)
-@click.option('--from-namespace', '-fn',
-              help='Origin Datastore namespace.',
-              required=True)
-@click.option('--to-project', '-tp',
-              help='Destination GCP Project.',
-              required=True)
-@click.option('--to-namespace', '-tn',
-              help='Destination Datastore namespace.',
-              required=True)
-@click.option('--skip-export',
-              help='Skip the export process from origin database.',
-              required=True,
-              is_flag=True,
-              default=False,
-              show_default=True)
-@click.option('--skip-import',
-              help='Skip the import process to destination database',
-              required=True,
-              is_flag=True,
-              default=False,
-              show_default=True)
-@click.option('--export-data-dir', help='Directory used to store exported data files.',
-              required=True,
-              default='./data',
-              show_default=True)
-@click.option('--kinds',
-              help='Comma separated list of Datastore Kinds to use.',
-              required=True)
-def migrate(from_project, from_namespace, to_project, to_namespace, skip_export, skip_import, export_data_dir, kinds):
-    """Migrate data from one namespace to another."""
-    kinds_list = kinds.split(',')
-
-    click.echo("Executing migration from '{}.{}' to '{}.{}'. Kinds: '{}'"
-               .format(from_project, from_namespace, to_project, to_namespace, kinds))
-    click.echo("Storing export data in '{}' directory."
-               .format(export_data_dir))
-
-    click.echo("Starting...")
-
-    if not skip_export:
-        with click.progressbar(kinds_list, label='Exporting',
-                               show_eta=False, item_show_func=show_progressbar_item) as bar:
-            for kind in bar:
-                # TODO: Implement export feature.
-                pass
-
-    if not skip_import:
-        with click.progressbar(kinds_list, label='Importing',
-                               show_eta=False, item_show_func=show_progressbar_item) as bar:
-            for kind in bar:
-                # TODO: Implement import feature.
-                pass
-
-    click.echo("Finished!")
-
-
-def show_progressbar_item(item):
-    return 'Kind: {}'.format(item)
+main.add_command(export)
+main.add_command(import_cmd)

gcdu/commands/export.py

@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+"""The export command."""
+import json
+import click
+
+from .utils import get_datastore_api, get_kinds_list, show_progressbar_item, partition_replace, save
+
+
+@click.command()
+@click.option('--project', '-p',
+              help='GCP project.',
+              required=True)
+@click.option('--namespace', '-n',
+              help='Datastore namespace.',
+              required=True)
+@click.option('--data-dir',
+              help='Directory to be used to store exported files.',
+              required=True,
+              default='./data',
+              show_default=True)
+@click.option('--project-placeholder', '-pp',
+              help='Placeholder value to replace the project value in exported files.',
+              required=True,
+              default='___PROJECT___',
+              show_default=True)
+@click.option('--namespace-placeholder', '-np',
+              help='Placeholder value to replace the namespace value in exported files.',
+              required=True,
+              default='___NAMESPACE___',
+              show_default=True)
+@click.option('--kinds', '-k',
+              help='Comma separated list of Datastore Kinds to export.',
+              required=True)
+def export(project, namespace, data_dir, project_placeholder, namespace_placeholder, kinds):
+    """Export data from database."""
+    kinds_list = get_kinds_list(kinds)
+    click.echo("Executing export. Project={}, Namespace={}, Kinds={}.".format(project, namespace, kinds_list))
+    with click.progressbar(kinds_list, label='Exporting', show_eta=True,
+                           item_show_func=show_progressbar_item) as bar:
+        for kind in bar:
+            execute_export(project, namespace, data_dir, project_placeholder, namespace_placeholder, kind)
+
+
+def execute_export(project, namespace, data_dir, project_placeholder, namespace_placeholder, kind):
+    datastore = get_datastore_api()
+
+    request_body = {
+        'partitionId': {
+            'projectId': project,
+            'namespaceId': namespace
+        },
+        'query': {
+            'kind': [
+                {
+                    'name': kind
+                }
+            ]
+        }
+    }
+
+    response = datastore.projects() \
+        .runQuery(projectId=project, body=request_body) \
+        .execute()
+
+    entities = extract_entities(response)
+    entities_json = json.dumps(entities)
+    entities_replaced_json = partition_replace(entities_json, project, project_placeholder, namespace,
+                                                       namespace_placeholder)
+    entities_replaced = json.loads(entities_replaced_json)
+
+    save(entities_replaced, kind, data_dir)
+
+
+def extract_entities(response):
+    entities = []
+    for entityResult in response.get('batch').get('entityResults'):
+        entity = entityResult.get('entity')
+        entities.append(entity)
+    return entities

gcdu/commands/import_cmd.py

@@ -0,0 +1,67 @@
+# -*- coding: utf-8 -*-
+"""The import command."""
+import json
+import click
+
+from .utils import get_datastore_api, get_kinds_list, show_progressbar_item, partition_replace, load
+
+
+@click.command('import')
+@click.option('--project', '-p',
+              help='GCP project.',
+              required=True)
+@click.option('--namespace', '-n',
+              help='Datastore namespace.',
+              required=True)
+@click.option('--data-dir',
+              help='Directory to be used to store exported files.',
+              required=True,
+              default='./data',
+              show_default=True)
+@click.option('--project-placeholder', '-pp',
+              help='Placeholder value to replace the project value in previously exported files.',
+              required=True,
+              default='___PROJECT___',
+              show_default=True)
+@click.option('--namespace-placeholder', '-np',
+              help='Placeholder value to replace the namespace value in previously exported files.',
+              required=True,
+              default='___NAMESPACE___',
+              show_default=True)
+@click.option('--kinds', '-k',
+              help='Comma separated list of Datastore Kinds to import.',
+              required=True)
+def import_cmd(project, namespace, data_dir, project_placeholder, namespace_placeholder, kinds):
+    """Import data to database using previously exported data as input."""
+    kinds_list = get_kinds_list(kinds)
+    click.echo("Executing import. Project={}, Namespace={}, Kinds={}.".format(project, namespace, kinds_list))
+    with click.progressbar(kinds_list, label='Importing', show_eta=True,
+                           item_show_func=show_progressbar_item) as bar:
+        for kind in bar:
+            execute_import(project, namespace, data_dir, project_placeholder, namespace_placeholder, kind)
+
+
+def execute_import(project, namespace, data_dir, project_placeholder, namespace_placeholder, kind):
+    datastore = get_datastore_api()
+
+    entities = load(kind, data_dir)
+    entities_json = json.dumps(entities)
+    entities_replaced_json = partition_replace(entities_json, project_placeholder, project,
+                                               namespace_placeholder, namespace)
+    entities_replaced = json.loads(entities_replaced_json)
+
+    inserts = []
+    for entity in entities_replaced:
+        insert = {
+            'insert': entity
+        }
+        inserts.append(insert)
+
+    request_body = {
+        "mutations": inserts,
+        "mode": "NON_TRANSACTIONAL"
+    }
+
+    datastore.projects() \
+        .commit(projectId=project, body=request_body) \
+        .execute()

gcdu/commands/utils.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+"""Utilities."""
+import io
+import os
+import json
+import googleapiclient.discovery
+
+
+def get_datastore_api():
+    return googleapiclient.discovery.build('datastore', 'v1')
+
+
+def get_kinds_list(kinds):
+    return kinds.split(',')
+
+
+def show_progressbar_item(item):
+    if item is None:
+        return 'Done!'
+    return "Kind: {}".format(item)
+
+
+def partition_replace(entities_json, from_project, to_project, from_namespace, to_namespace):
+    result = entities_json.replace('"projectId": "{}"'.format(from_project),
+                                   '"projectId": "{}"'.format(to_project))
+    result = result.replace('"namespaceId": "{}"'.format(from_namespace),
+                            '"namespaceId": "{}"'.format(to_namespace))
+    return result
+
+
+def save(entities, kind, data_dir):
+    if not os.path.exists(data_dir):
+        os.makedirs(data_dir)
+
+    with io.open('{}/{}.json'.format(data_dir, kind), 'w', encoding='utf-8') as export_file:
+        export_file.write(
+            json.dumps(entities, ensure_ascii=False, sort_keys=True, indent=2, separators=(',', ': ')))
+
+
+def load(kind, data_dir):
+    with io.open('{}/{}.json'.format(data_dir, kind), 'r', encoding='utf-8') as export_file:
+        entities = json.load(export_file)
+    return entities

requirements.txt

@@ -4,4 +4,4 @@ google-auth-httplib2==0.0.3
 click==6.7
 tox==2.9.1
 pytest==3.3.2
-pytest-cov==2.5.1
+pytest-cov==2.5.1

setup.py

@@ -108,8 +108,8 @@ def run(self):
         # Trove classifiers
         # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
         # 'Development Status :: 1 - Planning',
-        'Development Status :: 2 - Pre-Alpha',
-        # 'Development Status :: 3 - Alpha',
+        # 'Development Status :: 2 - Pre-Alpha',
+        'Development Status :: 3 - Alpha',
         # 'Development Status :: 4 - Beta',
         # 'Development Status :: 5 - Production/Stable',
         # 'Development Status :: 6 - Mature',