feat: support autoconversion of Entity to Key for purposes of delete & delete_multi (#123)

crwilcox · tseaver · web-flow · commit bf1dde60b2f4 · 2020-12-04T16:57:47.000-07:00
* feat: support autoconversion of Entity to Key for purposes of delete, delete_multi

* Update google/cloud/datastore/client.py

Co-authored-by: Tres Seaver &lt;tseaver@palladion.com&gt;

* test: update typing, use entity delete path

* fix: remove warning

* test: add unit test for delete multi with an entity

* fix: lint/black

Co-authored-by: Tres Seaver &lt;tseaver@palladion.com&gt;
diff --git a/google/cloud/datastore/client.py b/google/cloud/datastore/client.py
@@ -622,7 +622,8 @@ def delete(self, key, retry=None, timeout=None):
            The backend API does not make a distinction between a single key or
            multiple keys in a commit request.
 
-        :type key: :class:`google.cloud.datastore.key.Key`
+        :type key: :class:`google.cloud.datastore.key.Key`, :class:`google.cloud.datastore.entity.Entity`
+
         :param key: The key to be deleted from the datastore.
 
         :type retry: :class:`google.api_core.retry.Retry`
@@ -643,7 +644,7 @@ def delete(self, key, retry=None, timeout=None):
     def delete_multi(self, keys, retry=None, timeout=None):
         """Delete keys from the Cloud Datastore.
 
-        :type keys: list of :class:`google.cloud.datastore.key.Key`
+        :type keys: list of :class:`google.cloud.datastore.key.Key`, :class:`google.cloud.datastore.entity.Entity`
         :param keys: The keys to be deleted from the Datastore.
 
         :type retry: :class:`google.api_core.retry.Retry`
@@ -671,6 +672,9 @@ def delete_multi(self, keys, retry=None, timeout=None):
             current.begin()
 
         for key in keys:
+            if isinstance(key, Entity):
+                # If the key is in fact an Entity, the key can be extracted.
+                key = key.key
             current.delete(key)
 
         if not in_batch:
diff --git a/tests/system/test_system.py b/tests/system/test_system.py
@@ -72,9 +72,8 @@ def setUpModule():
 
 
 def tearDownModule():
-    keys = [entity.key for entity in Config.TO_DELETE]
     with Config.CLIENT.transaction():
-        Config.CLIENT.delete_multi(keys)
+        Config.CLIENT.delete_multi(Config.TO_DELETE)
 
 
 class TestDatastore(unittest.TestCase):
@@ -83,8 +82,7 @@ def setUp(self):
 
     def tearDown(self):
         with Config.CLIENT.transaction():
-            keys = [entity.key for entity in self.case_entities_to_delete]
-            Config.CLIENT.delete_multi(keys)
+            Config.CLIENT.delete_multi(self.case_entities_to_delete)
 
 
 class TestDatastoreAllocateIDs(TestDatastore):
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -961,6 +961,24 @@ def test_delete_multi_w_existing_transaction(self):
         self.assertEqual(mutated_key, key._key)
         client._datastore_api_internal.commit.assert_not_called()
 
+    def test_delete_multi_w_existing_transaction_entity(self):
+        from google.cloud.datastore.entity import Entity
+
+        creds = _make_credentials()
+        client = self._make_one(credentials=creds)
+        client._datastore_api_internal = _make_datastore_api()
+
+        key = _Key()
+        entity = Entity(key=key)
+
+        with _NoCommitTransaction(client) as CURR_XACT:
+            result = client.delete_multi([entity])
+
+        self.assertIsNone(result)
+        mutated_key = _mutated_pb(self, CURR_XACT.mutations, "delete")
+        self.assertEqual(mutated_key, key._key)
+        client._datastore_api_internal.commit.assert_not_called()
+
     def test_allocate_ids_w_partial_key(self):
         num_ids = 2
 
