Adding api.delete() and removing Key.delete().

Also:
- Changing api.get() back to accept only `keys` and
  returns only a list (a lot of headache for not much
  gain).
- Factored out behavior to extract shared dataset_id from
  a set of keys into _get_dataset_id_from_keys().
- Updated docstrings and other tests that rely on changed /
  removed methods.

See googleapis#518 for some context.

Author: dhermes

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

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

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
     """

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.

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

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

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

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