diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py index 876115b2dcaf..ff6048aebf00 100644 --- a/gcloud/datastore/batch.py +++ b/gcloud/datastore/batch.py @@ -63,10 +63,23 @@ class Batch(object): _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 def current(self): """Return the topmost batch / transaction, or None.""" @@ -202,18 +215,26 @@ def delete(self, key): self._add_delete_key_pb().CopyFrom(key_pb) def begin(self): - """No-op + """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:`gcloud.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): + 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. + 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 @@ -224,12 +245,26 @@ def commit(self): new_id = new_key_pb.path_element[-1].id entity.key = entity.key.completed_key(new_id) + 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. + """ + try: + self._commit() + finally: + self._status = self._FINISHED + def rollback(self): - """No-op + """Rolls back the current batch. + + Marks the batch as aborted (can't be used again). Overridden by :class:`gcloud.datastore.transaction.Transaction`. """ - pass + self._status = self._ABORTED def __enter__(self): self._client._push_batch(self) diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py index 657516d6a454..5267cb45b5b9 100644 --- a/gcloud/datastore/connection.py +++ b/gcloud/datastore/connection.py @@ -293,7 +293,7 @@ def begin_transaction(self, project): _datastore_pb2.BeginTransactionResponse) return response.transaction - def commit(self, project, commit_request, transaction_id): + def commit(self, project, request, transaction_id): """Commit mutations in context of current transation (if any). Maps the ``DatastoreService.Commit`` protobuf RPC. @@ -301,22 +301,23 @@ def commit(self, project, commit_request, transaction_id): :type project: string :param project: The project to which the transaction applies. - :type commit_request: :class:`._generated.datastore_pb2.CommitRequest` - :param commit_request: The protobuf with the mutations being committed. + :type request: :class:`._generated.datastore_pb2.CommitRequest` + :param request: The protobuf with the mutations being committed. :type transaction_id: string or None :param transaction_id: The transaction ID returned from :meth:`begin_transaction`. Non-transactional batches must pass ``None``. + .. note:: + + This method will mutate ``request`` before using it. + :rtype: tuple :returns': The pair of the number of index updates and a list of :class:`._generated.entity_pb2.Key` for each incomplete key that was completed in the commit. """ - request = _datastore_pb2.CommitRequest() - request.CopyFrom(commit_request) - if transaction_id: request.mode = _datastore_pb2.CommitRequest.TRANSACTIONAL request.transaction = transaction_id diff --git a/gcloud/datastore/test_batch.py b/gcloud/datastore/test_batch.py index b28d2af1a0ab..316070bcfeb1 100644 --- a/gcloud/datastore/test_batch.py +++ b/gcloud/datastore/test_batch.py @@ -37,6 +37,7 @@ def test_ctor(self): self.assertEqual(batch.connection, connection) self.assertEqual(batch.namespace, _NAMESPACE) self.assertTrue(batch._id is None) + self.assertEqual(batch._status, batch._INITIAL) self.assertTrue(isinstance(batch.mutations, datastore_pb2.Mutation)) self.assertEqual(batch._partial_key_entities, []) @@ -200,13 +201,39 @@ def test_delete_w_completed_key_w_prefixed_project(self): mutated_key = _mutated_pb(self, batch.mutations, 'delete') self.assertEqual(mutated_key, key._key) + def test_begin(self): + _PROJECT = 'PROJECT' + client = _Client(_PROJECT, None) + batch = self._makeOne(client) + self.assertEqual(batch._status, batch._INITIAL) + batch.begin() + self.assertEqual(batch._status, batch._IN_PROGRESS) + + def test_begin_fail(self): + _PROJECT = 'PROJECT' + client = _Client(_PROJECT, None) + batch = self._makeOne(client) + batch._status = batch._IN_PROGRESS + with self.assertRaises(ValueError): + batch.begin() + + def test_rollback(self): + _PROJECT = 'PROJECT' + client = _Client(_PROJECT, None) + batch = self._makeOne(client) + self.assertEqual(batch._status, batch._INITIAL) + batch.rollback() + self.assertEqual(batch._status, batch._ABORTED) + def test_commit(self): _PROJECT = 'PROJECT' connection = _Connection() client = _Client(_PROJECT, connection) batch = self._makeOne(client) + self.assertEqual(batch._status, batch._INITIAL) batch.commit() + self.assertEqual(batch._status, batch._FINISHED) self.assertEqual(connection._committed, [(_PROJECT, batch._commit_request, None)]) @@ -222,7 +249,9 @@ def test_commit_w_partial_key_entities(self): key._id = None batch._partial_key_entities.append(entity) + self.assertEqual(batch._status, batch._INITIAL) batch.commit() + self.assertEqual(batch._status, batch._FINISHED) self.assertEqual(connection._committed, [(_PROJECT, batch._commit_request, None)]) diff --git a/gcloud/datastore/transaction.py b/gcloud/datastore/transaction.py index 6c9426f14d9f..25c92729e9ca 100644 --- a/gcloud/datastore/transaction.py +++ b/gcloud/datastore/transaction.py @@ -90,22 +90,9 @@ class Transaction(Batch): :param client: the client used to connect to datastore. """ - _INITIAL = 0 - """Enum value for _INITIAL status of transaction.""" - - _IN_PROGRESS = 1 - """Enum value for _IN_PROGRESS status of transaction.""" - - _ABORTED = 2 - """Enum value for _ABORTED status of transaction.""" - - _FINISHED = 3 - """Enum value for _FINISHED status of transaction.""" - def __init__(self, client): super(Transaction, self).__init__(client) self._id = None - self._status = self._INITIAL @property def id(self): @@ -139,9 +126,7 @@ def begin(self): :raises: :class:`ValueError` if the transaction has already begun. """ - if self._status != self._INITIAL: - raise ValueError('Transaction already started previously.') - self._status = self._IN_PROGRESS + super(Transaction, self).begin() self._id = self.connection.begin_transaction(self.project) def rollback(self): @@ -155,7 +140,7 @@ def rollback(self): try: self.connection.rollback(self.project, self._id) finally: - self._status = self._ABORTED + super(Transaction, self).rollback() # Clear our own ID in case this gets accidentally reused. self._id = None @@ -173,6 +158,5 @@ def commit(self): try: super(Transaction, self).commit() finally: - self._status = self._FINISHED # Clear our own ID in case this gets accidentally reused. self._id = None