Skip to content

Commit

Permalink
feat: support autoconversion of Entity to Key for purposes of delete …
Browse files Browse the repository at this point in the history
…& 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 <tseaver@palladion.com>

* 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 <tseaver@palladion.com>
  • Loading branch information
crwilcox and tseaver committed Dec 4, 2020
1 parent 3fcefc4 commit bf1dde6
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 6 deletions.
8 changes: 6 additions & 2 deletions google/cloud/datastore/client.py
Expand Up @@ -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`
Expand All @@ -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`
Expand Down Expand Up @@ -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:
Expand Down
6 changes: 2 additions & 4 deletions tests/system/test_system.py
Expand Up @@ -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):
Expand All @@ -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):
Expand Down
18 changes: 18 additions & 0 deletions tests/unit/test_client.py
Expand Up @@ -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

Expand Down

0 comments on commit bf1dde6

Please sign in to comment.