diff --git a/gcloud/datastore/__init__.py b/gcloud/datastore/__init__.py index 2425398be7bf4..75a410fa20ccf 100644 --- a/gcloud/datastore/__init__.py +++ b/gcloud/datastore/__init__.py @@ -51,6 +51,7 @@ from gcloud import credentials from gcloud.datastore import _implicit_environ from gcloud.datastore.api import allocate_ids +from gcloud.datastore.api import delete from gcloud.datastore.api import get from gcloud.datastore.batch import Batch from gcloud.datastore.connection import Connection diff --git a/gcloud/datastore/api.py b/gcloud/datastore/api.py index 1d54c5d062d47..91a90c49f7b7b 100644 --- a/gcloud/datastore/api.py +++ b/gcloud/datastore/api.py @@ -18,9 +18,9 @@ Query objects rather than via protobufs. """ -import collections - from gcloud.datastore import _implicit_environ +from gcloud.datastore.batch import _BATCHES +from gcloud.datastore.batch import Batch from gcloud.datastore import helpers @@ -60,12 +60,33 @@ def _require_connection(connection=None): return connection -def get(key_or_keys, missing=None, deferred=None, connection=None): +def _get_dataset_id_from_keys(keys): + """Determines dataset ID from a list of keys. + + :type keys: list of :class:`gcloud.datastore.key.Key` + :param keys: The keys from the same dataset. + + :rtype: string + :returns: The dataset ID of the keys. + :raises: :class:`ValueError` if the key dataset IDs don't agree. + """ + dataset_id = keys[0].dataset_id + # Rather than creating a list or set of all dataset IDs, we iterate + # and check. We could allow the backend to check this for us if IDs + # with no prefix worked (GoogleCloudPlatform/google-cloud-datastore#59) + # or if we made sure that a prefix s~ or e~ was on each key. + for key in keys[1:]: + if key.dataset_id != dataset_id: + raise ValueError('All keys in get must be from the same dataset.') + + return dataset_id + + +def get(keys, missing=None, deferred=None, connection=None): """Retrieves entities, along with their attributes. - :type key_or_keys: list of :class:`gcloud.datastore.key.Key` or single - :class:`gcloud.datastore.key.Key` - :param key_or_keys: The key or keys to be retrieved from the datastore. + :type keys: list of :class:`gcloud.datastore.key.Key` + :param keys: The keys to be retrieved from the datastore. :type missing: an empty list or None. :param missing: If a list is passed, the key-only entities returned @@ -80,27 +101,14 @@ def get(key_or_keys, missing=None, deferred=None, connection=None): :type connection: :class:`gcloud.datastore.connection.Connection` :param connection: Optional. The connection used to connect to datastore. - :rtype: list of :class:`gcloud.datastore.entity.Entity`, single - :class:`gcloud.datastore.entity.Entity`, or ``NoneType`` - :returns: The requested entities, or single entity. + :rtype: list of :class:`gcloud.datastore.entity.Entity` + :returns: The requested entities. """ - if isinstance(key_or_keys, collections.Iterable): - keys = key_or_keys - else: - keys = [key_or_keys] - if not keys: return [] connection = _require_connection(connection) - dataset_id = keys[0].dataset_id - # Rather than creating a list or set of all dataset IDs, we iterate - # and check. We could allow the backend to check this for us if IDs - # with no prefix worked (GoogleCloudPlatform/google-cloud-datastore#59) - # or if we made sure that a prefix s~ or e~ was on each key. - for key in keys[1:]: - if key.dataset_id != dataset_id: - raise ValueError('All keys in get must be from the same dataset.') + dataset_id = _get_dataset_id_from_keys(keys) entity_pbs = connection.lookup( dataset_id=dataset_id, @@ -122,11 +130,33 @@ def get(key_or_keys, missing=None, deferred=None, connection=None): for entity_pb in entity_pbs: entities.append(helpers.entity_from_protobuf(entity_pb)) - if keys is key_or_keys: - return entities - else: - if entities: - return entities[0] + return entities + + +def delete(keys, connection=None): + """Delete the keys in the Cloud Datastore. + + :type keys: list of :class:`gcloud.datastore.key.Key` + :param keys: The keys to be deleted from the datastore. + + :type connection: :class:`gcloud.datastore.connection.Connection` + :param connection: Optional connection used to connect to datastore. + """ + if not keys: + return + + connection = connection or _implicit_environ.CONNECTION + + # We allow partial keys to attempt a delete, the backend will fail. + current = _BATCHES.top + in_batch = current is not None + if not in_batch: + dataset_id = _get_dataset_id_from_keys(keys) + current = Batch(dataset_id=dataset_id, connection=connection) + for key in keys: + current.delete(key) + if not in_batch: + current.commit() def allocate_ids(incomplete_key, num_ids, connection=None): diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py index e840d94ddc953..5dfaa91259a47 100644 --- a/gcloud/datastore/batch.py +++ b/gcloud/datastore/batch.py @@ -82,7 +82,7 @@ class Batch(object): operations and the delete operatiuon into the same mutation, and send them to the server in a single API request:: - >>> from gcloud import datastore + >>> from gcloud.datastore.batch import Batch >>> batch = Batch() >>> batch.put(entity1) >>> batch.put(entity2) @@ -102,7 +102,7 @@ class Batch(object): >>> from gcloud import datastore >>> dataset = datastore.get_dataset('dataset-id') - >>> with Batch as batch: + >>> with Batch() as batch: ... do_some_work(batch) ... raise Exception() # rolls back """ diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py index ec3253a3b2817..e74c6c2a29ec2 100644 --- a/gcloud/datastore/connection.py +++ b/gcloud/datastore/connection.py @@ -506,8 +506,7 @@ def delete_entities(self, dataset_id, key_pbs, mutation=None): This method deals only with :class:`gcloud.datastore.datastore_v1_pb2.Key` protobufs and not with any of the other abstractions. For example, it's used - under the hood in the - :meth:`gcloud.datastore.entity.Entity.delete` method. + under the hood in :func:`gcloud.datastore.api.delete`. :type dataset_id: string :param dataset_id: The ID of the dataset from which to delete the keys. diff --git a/gcloud/datastore/demo/demo.py b/gcloud/datastore/demo/demo.py index 6175a7adccedc..7a4c40eb3b3b4 100644 --- a/gcloud/datastore/demo/demo.py +++ b/gcloud/datastore/demo/demo.py @@ -34,7 +34,7 @@ print(datastore.get([toy.key])) # And we should be able to delete it... -toy.key.delete() +datastore.delete([toy.key]) # Since we deleted it, if we do another lookup it shouldn't be there again: print(datastore.get([toy.key])) @@ -48,10 +48,10 @@ (4567, 'Printer', 11), (5678, 'Printer', 12), (6789, 'Computer', 13)] -samples = [] +sample_keys = [] for id, name, age in SAMPLE_DATA: key = datastore.Key('Thing', id) - samples.append(key) + sample_keys.append(key) entity = datastore.Entity(key) entity['name'] = name entity['age'] = age @@ -72,7 +72,7 @@ print(list(query.fetch())) # Now delete them. -print([key.delete() for key in samples]) +datastore.delete(sample_keys) # You can also work inside a transaction. # (Check the official docs for explanations of what's happening here.) @@ -92,7 +92,7 @@ print('Committing the transaction...') # Now that the transaction is commited, let's delete the entities. -print(key.delete(), key2.delete()) +datastore.delete([key, key2]) # To rollback a transaction, just call .rollback() with datastore.Transaction() as t: @@ -116,4 +116,4 @@ print(thing.key) # This will be complete # Now let's delete the entity. -thing.key.delete() +datastore.delete([thing.key]) diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py index bf4d86efbd4e0..ce808ff9e5d2f 100644 --- a/gcloud/datastore/key.py +++ b/gcloud/datastore/key.py @@ -217,19 +217,6 @@ def to_protobuf(self): return key - def delete(self, connection=None): - """Delete the key in the Cloud Datastore. - - :type connection: :class:`gcloud.datastore.connection.Connection` - :param connection: Optional connection used to connect to datastore. - """ - # We allow partial keys to attempt a delete, the backend will fail. - connection = connection or _implicit_environ.CONNECTION - connection.delete_entities( - dataset_id=self.dataset_id, - key_pbs=[self.to_protobuf()], - ) - @property def is_partial(self): """Boolean indicating if the key has an ID (or name). diff --git a/gcloud/datastore/test_api.py b/gcloud/datastore/test_api.py index 77ad3122c6ec4..c8f05e087d534 100644 --- a/gcloud/datastore/test_api.py +++ b/gcloud/datastore/test_api.py @@ -101,23 +101,15 @@ def test_get_no_keys(self): results = self._callFUT([]) self.assertEqual(results, []) - def _miss_helper(self, expected_results, use_list=True): + def test_get_miss(self): from gcloud.datastore.key import Key from gcloud.datastore.test_connection import _Connection DATASET_ID = 'DATASET' connection = _Connection() key = Key('Kind', 1234, dataset_id=DATASET_ID) - if use_list: - key = [key] - results = self._callFUT(key, connection=connection) - self.assertEqual(results, expected_results) - - def test_get_miss(self): - self._miss_helper([], use_list=True) - - def test_get_miss_single_key(self): - self._miss_helper(None, use_list=False) + results = self._callFUT([key], connection=connection) + self.assertEqual(results, []) def test_get_miss_w_missing(self): from gcloud.datastore import datastore_v1_pb2 as datastore_pb @@ -248,33 +240,6 @@ def test_get_hit_multiple_keys_different_dataset(self): with self.assertRaises(ValueError): self._callFUT([key1, key2], connection=object()) - def test_get_hit_single_key(self): - from gcloud.datastore.key import Key - from gcloud.datastore.test_connection import _Connection - - DATASET_ID = 'DATASET' - KIND = 'Kind' - ID = 1234 - PATH = [{'kind': KIND, 'id': ID}] - - # Make a found entity pb to be returned from mock backend. - entity_pb = self._make_entity_pb(DATASET_ID, KIND, ID, - 'foo', 'Foo') - - # Make a connection to return the entity pb. - connection = _Connection(entity_pb) - - key = Key(KIND, ID, dataset_id=DATASET_ID) - result = self._callFUT(key, connection=connection) - new_key = result.key - - # Check the returned value is as expected. - self.assertFalse(new_key is key) - self.assertEqual(new_key.dataset_id, DATASET_ID) - self.assertEqual(new_key.path, PATH) - self.assertEqual(list(result), ['foo']) - self.assertEqual(result['foo'], 'Foo') - def test_get_implicit(self): from gcloud.datastore import _implicit_environ from gcloud.datastore.key import Key @@ -313,6 +278,98 @@ def test_get_implicit(self): self.assertEqual(result['foo'], 'Foo') +class Test_delete_function(unittest2.TestCase): + + def _callFUT(self, keys, connection=None): + from gcloud.datastore.api import delete + return delete(keys, connection=connection) + + def test_delete_no_batch(self): + from gcloud.datastore.test_batch import _Connection + from gcloud.datastore.test_batch import _Key + + # Build basic mocks needed to delete. + _DATASET = 'DATASET' + connection = _Connection() + key = _Key(_DATASET) + + result = self._callFUT([key], connection=connection) + self.assertEqual(result, None) + + def test_delete_existing_batch(self): + from gcloud._testing import _Monkey + from gcloud.datastore import api + from gcloud.datastore.batch import _Batches + from gcloud.datastore.batch import Batch + from gcloud.datastore.test_batch import _Connection + from gcloud.datastore.test_batch import _Key + + # Build basic mocks needed to delete. + _DATASET = 'DATASET' + connection = _Connection() + key = _Key(_DATASET) + + # Set up mock Batch on stack so we can check it is used. + _BATCHES = _Batches() + CURR_BATCH = Batch(dataset_id=_DATASET, connection=connection) + _BATCHES.push(CURR_BATCH) + + with _Monkey(api, _BATCHES=_BATCHES): + result = self._callFUT([key], connection=connection) + + self.assertEqual(result, None) + self.assertEqual( + connection._deleted, + [(_DATASET, [key._key], CURR_BATCH.mutation)]) + + def test_delete_implicit_connection(self): + from gcloud._testing import _Monkey + from gcloud.datastore import _implicit_environ + from gcloud.datastore import api + from gcloud.datastore.batch import _Batches + from gcloud.datastore.batch import Batch + from gcloud.datastore.test_batch import _Connection + from gcloud.datastore.test_batch import _Key + + # Build basic mocks needed to delete. + _DATASET = 'DATASET' + connection = _Connection() + key = _Key(_DATASET) + + # Set up mock Batch on stack so we can check it is used. + _BATCHES = _Batches() + + with _Monkey(_implicit_environ, CONNECTION=connection): + CURR_BATCH = Batch(dataset_id=_DATASET) + _BATCHES.push(CURR_BATCH) + with _Monkey(api, _BATCHES=_BATCHES): + result = self._callFUT([key]) + + self.assertEqual(result, None) + self.assertEqual( + connection._deleted, + [(_DATASET, [key._key], CURR_BATCH.mutation)]) + + def test_delete_no_keys(self): + from gcloud.datastore import _implicit_environ + + self.assertEqual(_implicit_environ.CONNECTION, None) + result = self._callFUT([]) + self.assertEqual(result, None) + + def test_delete_no_connection(self): + from gcloud.datastore import _implicit_environ + from gcloud.datastore.test_batch import _Key + + # Build basic mocks needed to delete. + _DATASET = 'DATASET' + key = _Key(_DATASET) + + self.assertEqual(_implicit_environ.CONNECTION, None) + with self.assertRaises(ValueError): + self._callFUT([key]) + + class Test_allocate_ids_function(unittest2.TestCase): def _callFUT(self, incomplete_key, num_ids, connection=None): diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py index d91c64555920c..e831488831955 100644 --- a/gcloud/datastore/test_connection.py +++ b/gcloud/datastore/test_connection.py @@ -1229,11 +1229,6 @@ def allocate_ids(self, dataset_id, key_pbs): num_pbs = len(key_pbs) return [_KeyProto(i) for i in range(num_pbs)] - def delete_entities(self, dataset_id, key_pbs): - self._called_dataset_id = dataset_id - self._called_key_pbs = key_pbs - return True - class _PathElementProto(object): diff --git a/gcloud/datastore/test_key.py b/gcloud/datastore/test_key.py index f95aa290eb2f8..704d8729b71b0 100644 --- a/gcloud/datastore/test_key.py +++ b/gcloud/datastore/test_key.py @@ -232,38 +232,6 @@ def test_to_protobuf_w_no_kind(self): pb = key.to_protobuf() self.assertFalse(pb.path_element[0].HasField('kind')) - def test_delete_explicit_connection(self): - from gcloud.datastore.test_connection import _Connection - - cnxn = _Connection() - key = self._makeOne('KIND', 1234, dataset_id=self._DEFAULT_DATASET) - result = key.delete(connection=cnxn) - self.assertEqual(result, None) - self.assertEqual(cnxn._called_dataset_id, self._DEFAULT_DATASET) - self.assertEqual(cnxn._called_key_pbs, [key.to_protobuf()]) - - def test_delete_implicit_connection(self): - from gcloud._testing import _Monkey - from gcloud.datastore import _implicit_environ - from gcloud.datastore.test_connection import _Connection - - cnxn = _Connection() - key = self._makeOne('KIND', 1234, dataset_id=self._DEFAULT_DATASET) - with _Monkey(_implicit_environ, CONNECTION=cnxn): - result = key.delete() - - self.assertEqual(result, None) - self.assertEqual(cnxn._called_dataset_id, self._DEFAULT_DATASET) - self.assertEqual(cnxn._called_key_pbs, [key.to_protobuf()]) - - def test_delete_no_connection(self): - from gcloud.datastore import _implicit_environ - - self.assertEqual(_implicit_environ.CONNECTION, None) - key = self._makeOne('KIND', 1234, dataset_id=self._DEFAULT_DATASET) - with self.assertRaises(AttributeError): - key.delete() - def test_is_partial_no_name_or_id(self): key = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET) self.assertTrue(key.is_partial) diff --git a/regression/clear_datastore.py b/regression/clear_datastore.py index 6923ee76bba64..88e76a7f6c241 100644 --- a/regression/clear_datastore.py +++ b/regression/clear_datastore.py @@ -51,19 +51,6 @@ def get_ancestors(entities): return list(set(key_roots)) -def delete_entities(entities): - if not entities: - return - - 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] - datastore.get_connection().delete_entities(dataset_id, key_pbs) - - def remove_kind(kind): results = [] @@ -87,10 +74,10 @@ def remove_kind(kind): if len(ancestors) > TRANSACTION_MAX_GROUPS: delete_outside_transaction = True else: - delete_entities(results) + datastore.delete([result.key for result in results]) if delete_outside_transaction: - delete_entities(results) + datastore.delete([result.key for result in results]) def remove_all_entities(): diff --git a/regression/datastore.py b/regression/datastore.py index 057cad270e4b6..7c73a20126838 100644 --- a/regression/datastore.py +++ b/regression/datastore.py @@ -33,8 +33,8 @@ def setUp(self): def tearDown(self): with datastore.Transaction(): - for entity in self.case_entities_to_delete: - entity.key.delete() + keys = [entity.key for entity in self.case_entities_to_delete] + datastore.delete(keys) class TestDatastoreAllocateIDs(TestDatastore): @@ -86,7 +86,7 @@ def _generic_test_post(self, name=None, key_id=None): self.assertEqual(entity.key.name, name) if key_id is not None: self.assertEqual(entity.key.id, key_id) - retrieved_entity = datastore.get(entity.key) + retrieved_entity, = datastore.get([entity.key]) # Check the keys are the same. self.assertEqual(retrieved_entity.key.path, entity.key.path) self.assertEqual(retrieved_entity.key.namespace, entity.key.namespace) @@ -341,13 +341,13 @@ def test_transaction(self): entity['url'] = u'www.google.com' with datastore.Transaction(): - retrieved_entity = datastore.get(entity.key) - if retrieved_entity is None: + results = datastore.get([entity.key]) + if len(results) == 0: entity.save() self.case_entities_to_delete.append(entity) # This will always return after the transaction. - retrieved_entity = datastore.get(entity.key) + retrieved_entity, = datastore.get([entity.key]) self.case_entities_to_delete.append(retrieved_entity) retrieved_dict = dict(retrieved_entity.items()) entity_dict = dict(entity.items())