# Copyright 2014 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Create / interact with a batch of updates / deletes.
Batches provide the ability to execute multiple operations
in a single request to the Cloud Datastore API.
See
https://cloud.google.com/datastore/docs/concepts/entities#Datastore_Batch_operations
"""
from google.cloud.datastore import helpers
from google.cloud.datastore._generated import datastore_pb2 as _datastore_pb2
[docs]class Batch(object):
"""An abstraction representing a collected group of updates / deletes.
Used to build up a bulk mutation.
For example, the following snippet of code will put the two ``save``
operations and the ``delete`` operation into the same mutation, and send
them to the server in a single API request::
>>> from google.cloud import datastore
>>> client = datastore.Client()
>>> batch = client.batch()
>>> batch.put(entity1)
>>> batch.put(entity2)
>>> batch.delete(key3)
>>> batch.commit()
You can also use a batch as a context manager, in which case
:meth:`commit` will be called automatically if its block exits without
raising an exception::
>>> with batch:
... batch.put(entity1)
... batch.put(entity2)
... batch.delete(key3)
By default, no updates will be sent if the block exits with an error::
>>> with batch:
... do_some_work(batch)
... raise Exception() # rolls back
:type client: :class:`google.cloud.datastore.client.Client`
:param client: The client used to connect to datastore.
"""
_id = None # "protected" attribute, always None for non-transactions
_INITIAL = 0
"""Enum value for _INITIAL status of batch/transaction."""
_IN_PROGRESS = 1
"""Enum value for _IN_PROGRESS status of batch/transaction."""
_ABORTED = 2
"""Enum value for _ABORTED status of batch/transaction."""
_FINISHED = 3
"""Enum value for _FINISHED status of batch/transaction."""
def __init__(self, client):
self._client = client
self._commit_request = _datastore_pb2.CommitRequest()
self._partial_key_entities = []
self._status = self._INITIAL
[docs] def current(self):
"""Return the topmost batch / transaction, or None."""
return self._client.current_batch
@property
def project(self):
"""Getter for project in which the batch will run.
:rtype: :class:`str`
:returns: The project in which the batch will run.
"""
return self._client.project
@property
def namespace(self):
"""Getter for namespace in which the batch will run.
:rtype: :class:`str`
:returns: The namespace in which the batch will run.
"""
return self._client.namespace
@property
def connection(self):
"""Getter for connection over which the batch will run.
:rtype: :class:`google.cloud.datastore.connection.Connection`
:returns: The connection over which the batch will run.
"""
return self._client.connection
def _add_partial_key_entity_pb(self):
"""Adds a new mutation for an entity with a partial key.
:rtype: :class:`google.cloud.datastore._generated.entity_pb2.Entity`
:returns: The newly created entity protobuf that will be
updated and sent with a commit.
"""
new_mutation = self.mutations.add()
return new_mutation.insert
def _add_complete_key_entity_pb(self):
"""Adds a new mutation for an entity with a completed key.
:rtype: :class:`google.cloud.datastore._generated.entity_pb2.Entity`
:returns: The newly created entity protobuf that will be
updated and sent with a commit.
"""
# We use ``upsert`` for entities with completed keys, rather than
# ``insert`` or ``update``, in order not to create race conditions
# based on prior existence / removal of the entity.
new_mutation = self.mutations.add()
return new_mutation.upsert
def _add_delete_key_pb(self):
"""Adds a new mutation for a key to be deleted.
:rtype: :class:`google.cloud.datastore._generated.entity_pb2.Key`
:returns: The newly created key protobuf that will be
deleted when sent with a commit.
"""
new_mutation = self.mutations.add()
return new_mutation.delete
@property
def mutations(self):
"""Getter for the changes accumulated by this batch.
Every batch is committed with a single commit request containing all
the work to be done as mutations. Inside a batch, calling :meth:`put`
with an entity, or :meth:`delete` with a key, builds up the request by
adding a new mutation. This getter returns the protobuf that has been
built-up so far.
:rtype: iterable
:returns: The list of :class:`._generated.datastore_pb2.Mutation`
protobufs to be sent in the commit request.
"""
return self._commit_request.mutations
[docs] def put(self, entity):
"""Remember an entity's state to be saved during :meth:`commit`.
.. note::
Any existing properties for the entity will be replaced by those
currently set on this instance. Already-stored properties which do
not correspond to keys set on this instance will be removed from
the datastore.
.. note::
Property values which are "text" ('unicode' in Python2, 'str' in
Python3) map to 'string_value' in the datastore; values which are
"bytes" ('str' in Python2, 'bytes' in Python3) map to 'blob_value'.
When an entity has a partial key, calling :meth:`commit` sends it as
an ``insert`` mutation and the key is completed. On return,
the key for the ``entity`` passed in is updated to match the key ID
assigned by the server.
:type entity: :class:`google.cloud.datastore.entity.Entity`
:param entity: the entity to be saved.
:raises: :class:`~exceptions.ValueError` if the batch is not in
progress, if entity has no key assigned, or if the key's
``project`` does not match ours.
"""
if self._status != self._IN_PROGRESS:
raise ValueError('Batch must be in progress to put()')
if entity.key is None:
raise ValueError("Entity must have a key")
if self.project != entity.key.project:
raise ValueError("Key must be from same project as batch")
if entity.key.is_partial:
entity_pb = self._add_partial_key_entity_pb()
self._partial_key_entities.append(entity)
else:
entity_pb = self._add_complete_key_entity_pb()
_assign_entity_to_pb(entity_pb, entity)
[docs] def delete(self, key):
"""Remember a key to be deleted during :meth:`commit`.
:type key: :class:`google.cloud.datastore.key.Key`
:param key: the key to be deleted.
:raises: :class:`~exceptions.ValueError` if the batch is not in
progress, if key is not complete, or if the key's
``project`` does not match ours.
"""
if self._status != self._IN_PROGRESS:
raise ValueError('Batch must be in progress to delete()')
if key.is_partial:
raise ValueError("Key must be complete")
if self.project != key.project:
raise ValueError("Key must be from same project as batch")
key_pb = key.to_protobuf()
self._add_delete_key_pb().CopyFrom(key_pb)
[docs] def begin(self):
"""Begins a batch.
This method is called automatically when entering a with
statement, however it can be called explicitly if you don't want
to use a context manager.
Overridden by :class:`google.cloud.datastore.transaction.Transaction`.
:raises: :class:`ValueError` if the batch has already begun.
"""
if self._status != self._INITIAL:
raise ValueError('Batch already started previously.')
self._status = self._IN_PROGRESS
def _commit(self):
"""Commits the batch.
This is called by :meth:`commit`.
"""
# NOTE: ``self._commit_request`` will be modified.
_, updated_keys = self.connection.commit(
self.project, self._commit_request, self._id)
# If the back-end returns without error, we are guaranteed that
# :meth:`Connection.commit` will return keys that match (length and
# order) directly ``_partial_key_entities``.
for new_key_pb, entity in zip(updated_keys,
self._partial_key_entities):
new_id = new_key_pb.path[-1].id
entity.key = entity.key.completed_key(new_id)
[docs] def commit(self):
"""Commits the batch.
This is called automatically upon exiting a with statement,
however it can be called explicitly if you don't want to use a
context manager.
:raises: :class:`~exceptions.ValueError` if the batch is not
in progress.
"""
if self._status != self._IN_PROGRESS:
raise ValueError('Batch must be in progress to commit()')
try:
self._commit()
finally:
self._status = self._FINISHED
[docs] def rollback(self):
"""Rolls back the current batch.
Marks the batch as aborted (can't be used again).
Overridden by :class:`google.cloud.datastore.transaction.Transaction`.
:raises: :class:`~exceptions.ValueError` if the batch is not
in progress.
"""
if self._status != self._IN_PROGRESS:
raise ValueError('Batch must be in progress to rollback()')
self._status = self._ABORTED
def __enter__(self):
self.begin()
# NOTE: We make sure begin() succeeds before pushing onto the stack.
self._client._push_batch(self)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
try:
if exc_type is None:
self.commit()
else:
self.rollback()
finally:
self._client._pop_batch()
def _assign_entity_to_pb(entity_pb, entity):
"""Copy ``entity`` into ``entity_pb``.
Helper method for ``Batch.put``.
:type entity_pb: :class:`._generated.entity_pb2.Entity`
:param entity_pb: The entity owned by a mutation.
:type entity: :class:`google.cloud.datastore.entity.Entity`
:param entity: The entity being updated within the batch / transaction.
"""
bare_entity_pb = helpers.entity_to_protobuf(entity)
bare_entity_pb.key.CopyFrom(bare_entity_pb.key)
entity_pb.CopyFrom(bare_entity_pb)