Update regression scripts to work with new API.

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__':

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()
 
 

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'],