diff --git a/regression/clear_datastore.py b/regression/clear_datastore.py index bda94522ecba..3a8223231a34 100644 --- a/regression/clear_datastore.py +++ b/regression/clear_datastore.py @@ -15,13 +15,18 @@ """Script to populate datastore with regression test data.""" -# This assumes the command is being run via tox hence the -# repository root is the current directory. +from gcloud import datastore +from gcloud.datastore import _implicit_environ from gcloud.datastore.query import Query -from regression import regression_utils +from gcloud.datastore.transaction import Transaction from six.moves import input +datastore._DATASET_ENV_VAR_NAME = 'GCLOUD_TESTS_DATASET_ID' +datastore.set_default_dataset_id() +datastore.set_default_connection() + + FETCH_MAX = 20 ALL_KINDS = [ 'Character', @@ -33,10 +38,9 @@ TRANSACTION_MAX_GROUPS = 5 -def fetch_keys(dataset, kind, fetch_max=FETCH_MAX, query=None, cursor=None): +def fetch_keys(kind, fetch_max=FETCH_MAX, query=None, cursor=None): if query is None: - query = Query( - dataset_id=dataset.id(), kind=kind, projection=['__key__']) + query = Query(kind=kind, projection=['__key__']) iterator = query.fetch(limit=fetch_max, start_cursor=cursor) @@ -46,37 +50,39 @@ def fetch_keys(dataset, kind, fetch_max=FETCH_MAX, query=None, cursor=None): def get_ancestors(entities): # NOTE: A key will always have at least one path element. - key_roots = [entity.key().path[0] for entity in entities] - # Turn into hashable type so we can use set to get unique roots. - # Also sorted the items() to ensure uniqueness. - key_roots = [tuple(sorted(root.items())) for root in key_roots] - # Cast back to dictionary. - return [dict(root) for root in set(key_roots)] + key_roots = [entity.key.flat_path[:2] for entity in entities] + # Return the unique roots. + return list(set(key_roots)) -def delete_entities(dataset, entities): - dataset_id = dataset.id() - connection = dataset.connection() +def delete_entities(entities): + if not entities: + return - key_pbs = [entity.key().to_protobuf() for entity in entities] - connection.delete_entities(dataset_id, key_pbs) + dataset_ids = set(entity.key.dataset_id for entity in entities) + if len(dataset_ids) != 1: + raise ValueError('Expected a unique dataset ID.') + dataset_id = dataset_ids.pop() + key_pbs = [entity.key.to_protobuf() for entity in entities] + _implicit_environ.CONNECTION.delete_entities(dataset_id, key_pbs) -def remove_kind(dataset, kind): - delete_outside_transaction = False - with dataset.transaction(): - results = [] - query, curr_results, cursor = fetch_keys(dataset, kind) +def remove_kind(kind): + results = [] + + query, curr_results, cursor = fetch_keys(kind) + results.extend(curr_results) + while curr_results: + query, curr_results, cursor = fetch_keys(kind, query=query, + cursor=cursor) results.extend(curr_results) - while curr_results: - query, curr_results, cursor = fetch_keys( - dataset, kind, query=query, cursor=cursor) - results.extend(curr_results) - if not results: - return + if not results: + return + delete_outside_transaction = False + with Transaction(): # Now that we have all results, we seek to delete. print('Deleting keys:') print(results) @@ -85,10 +91,10 @@ def remove_kind(dataset, kind): if len(ancestors) > TRANSACTION_MAX_GROUPS: delete_outside_transaction = True else: - delete_entities(dataset, results) + delete_entities(results) if delete_outside_transaction: - delete_entities(dataset, results) + delete_entities(results) def remove_all_entities(): @@ -99,9 +105,8 @@ def remove_all_entities(): print('Doing nothing.') return - dataset = regression_utils.get_dataset() for kind in ALL_KINDS: - remove_kind(dataset, kind) + remove_kind(kind) if __name__ == '__main__': diff --git a/regression/datastore.py b/regression/datastore.py index a6e78449b549..78b5597125ca 100644 --- a/regression/datastore.py +++ b/regression/datastore.py @@ -13,7 +13,6 @@ # limitations under the License. import datetime -import os import pytz import unittest2 @@ -27,8 +26,8 @@ from regression import populate_datastore -DATASET_ID = os.getenv('GCLOUD_TESTS_DATASET_ID') -datastore.set_default_dataset_id(dataset_id=DATASET_ID) +datastore._DATASET_ENV_VAR_NAME = 'GCLOUD_TESTS_DATASET_ID' +datastore.set_default_dataset_id() datastore.set_default_connection() diff --git a/regression/populate_datastore.py b/regression/populate_datastore.py index c773fd70b40f..c04f39ce155e 100644 --- a/regression/populate_datastore.py +++ b/regression/populate_datastore.py @@ -17,9 +17,14 @@ from six.moves import zip from gcloud import datastore -# This assumes the command is being run via tox hence the -# repository root is the current directory. -from regression import regression_utils +from gcloud.datastore.entity import Entity +from gcloud.datastore.key import Key +from gcloud.datastore.transaction import Transaction + + +datastore._DATASET_ENV_VAR_NAME = 'GCLOUD_TESTS_DATASET_ID' +datastore.set_default_dataset_id() +datastore.set_default_connection() ANCESTOR = ('Book', 'GoT') @@ -81,14 +86,12 @@ def add_characters(): - dataset = regression_utils.get_dataset() - with dataset.transaction(): + with Transaction(): for key_path, character in zip(KEY_PATHS, CHARACTERS): if key_path[-1] != character['name']: raise ValueError(('Character and key don\'t agree', key_path, character)) - key = datastore.key.Key(*key_path, dataset_id=dataset.id()) - entity = datastore.entity.Entity(dataset=dataset).key(key) + entity = Entity(key=Key(*key_path)) entity.update(character) entity.save() print('Adding Character %s %s' % (character['name'],