From bf1dde60b2f42e939c7dfa4a5228c3f41d565ece Mon Sep 17 00:00:00 2001 From: Christopher Wilcox Date: Fri, 4 Dec 2020 15:57:47 -0800 Subject: [PATCH] feat: support autoconversion of Entity to Key for purposes of delete & delete_multi (#123) * feat: support autoconversion of Entity to Key for purposes of delete, delete_multi * Update google/cloud/datastore/client.py Co-authored-by: Tres Seaver * 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 --- google/cloud/datastore/client.py | 8 ++++++-- tests/system/test_system.py | 6 ++---- tests/unit/test_client.py | 18 ++++++++++++++++++ 3 files changed, 26 insertions(+), 6 deletions(-) diff --git a/google/cloud/datastore/client.py b/google/cloud/datastore/client.py index 24b53b54..28d968ce 100644 --- 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 index 89b0a395..a91b99ae 100644 --- 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 index 55a45c7f..3c75a5fb 100644 --- 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