|
| 1 | +# Copyright 2014 Google Inc. All rights reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Create / interact with a batch of updates / deletes.""" |
| 16 | + |
| 17 | +from gcloud.datastore import _implicit_environ |
| 18 | +from gcloud.datastore import datastore_v1_pb2 as datastore_pb |
| 19 | + |
| 20 | + |
| 21 | +class Batch(object): |
| 22 | + """An abstraction representing a collected group of updates / deletes. |
| 23 | +
|
| 24 | + Used to build up a bulk mutuation. |
| 25 | +
|
| 26 | + For example, the following snippet of code will put the two ``save`` |
| 27 | + operations and the delete operatiuon into the same mutation, and send |
| 28 | + them to the server in a single API request:: |
| 29 | +
|
| 30 | + >>> from gcloud import datastore |
| 31 | + >>> batch = Batch() |
| 32 | + >>> batch.put(entity1) |
| 33 | + >>> batch.put(entity2) |
| 34 | + >>> batch.delete(key3) |
| 35 | + >>> batch.commit() |
| 36 | +
|
| 37 | + You can also use a batch as a context manager, in which case the |
| 38 | + ``commit`` will be called automatically if its block exits without |
| 39 | + raising an exception:: |
| 40 | +
|
| 41 | + >>> with Batch() as batch: |
| 42 | + ... batch.put(entity1) |
| 43 | + ... batch.put(entity2) |
| 44 | + ... batch.delete(key3) |
| 45 | +
|
| 46 | + By default, no updates will be sent if the block exits with an error:: |
| 47 | +
|
| 48 | + >>> from gcloud import datastore |
| 49 | + >>> dataset = datastore.get_dataset('dataset-id') |
| 50 | + >>> with Batch as batch: |
| 51 | + ... do_some_work(batch) |
| 52 | + ... raise Exception() # rolls back |
| 53 | + """ |
| 54 | + |
| 55 | + def __init__(self, dataset_id=None, connection=None): |
| 56 | + """ Construct a batch. |
| 57 | +
|
| 58 | + :type dataset_id: :class:`str`. |
| 59 | + :param dataset_id: The ID of the dataset. |
| 60 | +
|
| 61 | + :type connection: :class:`gcloud.datastore.connection.Connection` |
| 62 | + :param connection: The connection used to connect to datastore. |
| 63 | +
|
| 64 | + :raises: :class:`ValueError` if either a connection or dataset ID |
| 65 | + are not set. |
| 66 | + """ |
| 67 | + self._connection = connection or _implicit_environ.CONNECTION |
| 68 | + self._dataset_id = dataset_id or _implicit_environ.DATASET_ID |
| 69 | + |
| 70 | + if self._connection is None or self._dataset_id is None: |
| 71 | + raise ValueError('A batch must have a connection and ' |
| 72 | + 'a dataset ID set.') |
| 73 | + |
| 74 | + self._mutation = datastore_pb.Mutation() |
| 75 | + |
| 76 | + @property |
| 77 | + def dataset_id(self): |
| 78 | + """Getter for dataset ID in which the batch will run. |
| 79 | +
|
| 80 | + :rtype: :class:`str` |
| 81 | + :returns: The dataset ID in which the batch will run. |
| 82 | + """ |
| 83 | + return self._dataset_id |
| 84 | + |
| 85 | + @property |
| 86 | + def connection(self): |
| 87 | + """Getter for connection over which the batch will run. |
| 88 | +
|
| 89 | + :rtype: :class:`gcloud.datastore.connection.Connection` |
| 90 | + :returns: The connection over which the batch will run. |
| 91 | + """ |
| 92 | + return self._connection |
| 93 | + |
| 94 | + @property |
| 95 | + def mutation(self): |
| 96 | + """Getter for the current mutation. |
| 97 | +
|
| 98 | + Every batch is committed with a single Mutation |
| 99 | + representing the 'work' to be done as part of the batch. |
| 100 | + Inside a batch, calling ``batch.put()`` with an entity, or |
| 101 | + ``batch.delete`` with a key, builds up the mutation. |
| 102 | + This getter returns the Mutation protobuf that |
| 103 | + has been built-up so far. |
| 104 | +
|
| 105 | + :rtype: :class:`gcloud.datastore.datastore_v1_pb2.Mutation` |
| 106 | + :returns: The Mutation protobuf to be sent in the commit request. |
| 107 | + """ |
| 108 | + return self._mutation |
| 109 | + |
| 110 | + def put(self, entity): |
| 111 | + """Remember an entity's state to be saved during ``commit``. |
| 112 | +
|
| 113 | + .. note:: |
| 114 | + Any existing properties for the entity will be replaced by those |
| 115 | + currently set on this instance. Already-stored properties which do |
| 116 | + not correspond to keys set on this instance will be removed from |
| 117 | + the datastore. |
| 118 | +
|
| 119 | + .. note:: |
| 120 | + Property values which are "text" ('unicode' in Python2, 'str' in |
| 121 | + Python3) map to 'string_value' in the datastore; values which are |
| 122 | + "bytes" ('str' in Python2, 'bytes' in Python3) map to 'blob_value'. |
| 123 | +
|
| 124 | + :type entity: :class:`gcloud.datastore.entity.Entity` |
| 125 | + :param entity: the entity to be saved. |
| 126 | +
|
| 127 | + :raises: ValueError if entity has no key assigned. |
| 128 | + """ |
| 129 | + if entity.key is None: |
| 130 | + raise ValueError("Entity must have a key") |
| 131 | + |
| 132 | + key_pb = entity.key.to_protobuf() |
| 133 | + properties = dict(entity) |
| 134 | + exclude = tuple(entity.exclude_from_indexes) |
| 135 | + |
| 136 | + self.connection.save_entity( |
| 137 | + self.dataset_id, key_pb, properties, |
| 138 | + exclude_from_indexes=exclude, mutation=self.mutation) |
| 139 | + |
| 140 | + def delete(self, key): |
| 141 | + """Remember a key to be deleted durring ``commit``. |
| 142 | +
|
| 143 | + :type key: :class:`gcloud.datastore.key.Key` |
| 144 | + :param key: the key to be deleted. |
| 145 | +
|
| 146 | + :raises: ValueError if key is not complete. |
| 147 | + """ |
| 148 | + if key.is_partial: |
| 149 | + raise ValueError("Key must be complete") |
| 150 | + |
| 151 | + key_pb = key.to_protobuf() |
| 152 | + self.connection.delete_entities( |
| 153 | + self.dataset_id, [key_pb], mutation=self.mutation) |
| 154 | + |
| 155 | + def commit(self): |
| 156 | + """Commits the batch. |
| 157 | +
|
| 158 | + This is called automatically upon exiting a with statement, |
| 159 | + however it can be called explicitly if you don't want to use a |
| 160 | + context manager. |
| 161 | + """ |
| 162 | + self.connection.commit(self._dataset_id, self.mutation) |
| 163 | + |
| 164 | + def __enter__(self): |
| 165 | + return self |
| 166 | + |
| 167 | + def __exit__(self, exc_type, exc_val, exc_tb): |
| 168 | + if exc_type is None: |
| 169 | + self.commit() |
0 commit comments