Repository: googleapis/google-cloud-python

Issue:
Implement allocate_ids for datastore
I discovered this when checking out the `gcloud-node` [regression tests](https://github.com/GoogleCloudPlatform/gcloud-node/blob/5865e36491ff6cbe2a47a91e576ab92a71dec7b1/regression/datastore.js#L31).

Implement allocate_ids for datastore
I discovered this when checking out the `gcloud-node` [regression tests](https://github.com/GoogleCloudPlatform/gcloud-node/blob/5865e36491ff6cbe2a47a91e576ab92a71dec7b1/regression/datastore.js#L31).



Discussion:
I'm not sure what that test in the node code is exercising, but it seems like it might be the following:

  https://github.com/GoogleCloudPlatform/gcloud-python/blob/41f97e187173e9f36930258bfd9e54048d22fe97/gcloud/datastore/transaction.py#L237-L240

where the transaction updates partial keys after committing, using the results from the commit response.

@tseaver is [this endpoint of the API](https://cloud.google.com/datastore/docs/apis/v1beta2/datasets/allocateIds)

I found the [protobuf request and response](https://github.com/GoogleCloudPlatform/google-cloud-datastore/blob/v1beta2-rev1-2.1.0/proto/datastore_v1.proto#L555-L568)

I'm not sure what that test in the node code is exercising, but it seems like it might be the following:

  https://github.com/GoogleCloudPlatform/gcloud-python/blob/41f97e187173e9f36930258bfd9e54048d22fe97/gcloud/datastore/transaction.py#L237-L240

where the transaction updates partial keys after committing, using the results from the commit response.

@tseaver is [this endpoint of the API](https://cloud.google.com/datastore/docs/apis/v1beta2/datasets/allocateIds)

I found the [protobuf request and response](https://github.com/GoogleCloudPlatform/google-cloud-datastore/blob/v1beta2-rev1-2.1.0/proto/datastore_v1.proto#L555-L568)


Resolving patch:
diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py
--- a/gcloud/datastore/connection.py
+++ b/gcloud/datastore/connection.py
@@ -153,52 +153,77 @@ def dataset(self, *args, **kwargs):
         kwargs['connection'] = self
         return Dataset(*args, **kwargs)
 
-    def begin_transaction(self, dataset_id, serializable=False):
-        """Begin a transaction.
-
-        :type dataset_id: string
-        :param dataset_id: The dataset over which to execute the transaction.
-        """
+    def lookup(self, dataset_id, key_pbs):
+        """Lookup keys from a dataset in the Cloud Datastore.
 
-        if self.transaction():
-            raise ValueError('Cannot start a transaction with another already '
-                             'in progress.')
+        Maps the ``DatastoreService.Lookup`` protobuf RPC.
 
-        request = datastore_pb.BeginTransactionRequest()
+        This method deals only with protobufs
+        (:class:`gcloud.datastore.datastore_v1_pb2.Key`
+        and
+        :class:`gcloud.datastore.datastore_v1_pb2.Entity`)
+        and is used under the hood for methods like
+        :func:`gcloud.datastore.dataset.Dataset.get_entity`:
 
-        if serializable:
-            request.isolation_level = (
-                datastore_pb.BeginTransactionRequest.SERIALIZABLE)
-        else:
-            request.isolation_level = (
-                datastore_pb.BeginTransactionRequest.SNAPSHOT)
+        >>> from gcloud import datastore
+        >>> from gcloud.datastore.key import Key
+        >>> connection = datastore.get_connection(email, key_path)
+        >>> dataset = connection.dataset('dataset-id')
+        >>> key = Key(dataset=dataset).kind('MyKind').id(1234)
 
-        response = self._rpc(dataset_id, 'beginTransaction', request,
-                             datastore_pb.BeginTransactionResponse)
+        Using the :class:`gcloud.datastore.dataset.Dataset` helper:
 
-        return response.transaction
+        >>> dataset.get_entity(key)
+        <Entity object>
 
-    def rollback_transaction(self, dataset_id):
-        """Rollback the connection's existing transaction.
+        Using the ``connection`` class directly:
 
-        Raises a ``ValueError``
-        if the connection isn't currently in a transaction.
+        >>> connection.lookup('dataset-id', key.to_protobuf())
+        <Entity protobuf>
 
         :type dataset_id: string
-        :param dataset_id: The dataset to which the transaction belongs.
+        :param dataset_id: The dataset to look up the keys.
+
+        :type key_pbs: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
+                       (or a single Key)
+        :param key_pbs: The key (or keys) to retrieve from the datastore.
+
+        :rtype: list of :class:`gcloud.datastore.datastore_v1_pb2.Entity`
+                (or a single Entity)
+        :returns: The entities corresponding to the keys provided.
+                  If a single key was provided and no results matched,
+                  this will return None.
+                  If multiple keys were provided and no results matched,
+                  this will return an empty list.
         """
-        if not self.transaction() or not self.transaction().id():
-            raise ValueError('No transaction to rollback.')
+        lookup_request = datastore_pb.LookupRequest()
 
-        request = datastore_pb.RollbackRequest()
-        request.transaction = self.transaction().id()
-        # Nothing to do with this response, so just execute the method.
-        self._rpc(dataset_id, 'rollback', request,
-                  datastore_pb.RollbackResponse)
+        single_key = isinstance(key_pbs, datastore_pb.Key)
+
+        if single_key:
+            key_pbs = [key_pbs]
+
+        for key_pb in key_pbs:
+            lookup_request.key.add().CopyFrom(key_pb)
+
+        lookup_response = self._rpc(dataset_id, 'lookup', lookup_request,
+                                    datastore_pb.LookupResponse)
+
+        results = [result.entity for result in lookup_response.found]
+
+        if single_key:
+            if results:
+                return results[0]
+            else:
+                return None
+
+        return results
 
     def run_query(self, dataset_id, query_pb, namespace=None):
         """Run a query on the Cloud Datastore.
 
+        Maps the ``DatastoreService.RunQuery`` protobuf RPC.
+
         Given a Query protobuf,
         sends a ``runQuery`` request to the Cloud Datastore API
         and returns a list of entity protobufs matching the query.
@@ -250,73 +275,38 @@ def run_query(self, dataset_id, query_pb, namespace=None):
             response.batch.skipped_results,
         )
 
-    def lookup(self, dataset_id, key_pbs):
-        """Lookup keys from a dataset in the Cloud Datastore.
-
-        This method deals only with protobufs
-        (:class:`gcloud.datastore.datastore_v1_pb2.Key`
-        and
-        :class:`gcloud.datastore.datastore_v1_pb2.Entity`)
-        and is used under the hood for methods like
-        :func:`gcloud.datastore.dataset.Dataset.get_entity`:
-
-        >>> from gcloud import datastore
-        >>> from gcloud.datastore.key import Key
-        >>> connection = datastore.get_connection(email, key_path)
-        >>> dataset = connection.dataset('dataset-id')
-        >>> key = Key(dataset=dataset).kind('MyKind').id(1234)
-
-        Using the :class:`gcloud.datastore.dataset.Dataset` helper:
-
-        >>> dataset.get_entity(key)
-        <Entity object>
-
-        Using the ``connection`` class directly:
+    def begin_transaction(self, dataset_id, serializable=False):
+        """Begin a transaction.
 
-        >>> connection.lookup('dataset-id', key.to_protobuf())
-        <Entity protobuf>
+        Maps the ``DatastoreService.BeginTransaction`` protobuf RPC.
 
         :type dataset_id: string
-        :param dataset_id: The dataset to look up the keys.
-
-        :type key_pbs: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
-                       (or a single Key)
-        :param key_pbs: The key (or keys) to retrieve from the datastore.
-
-        :rtype: list of :class:`gcloud.datastore.datastore_v1_pb2.Entity`
-                (or a single Entity)
-        :returns: The entities corresponding to the keys provided.
-                  If a single key was provided and no results matched,
-                  this will return None.
-                  If multiple keys were provided and no results matched,
-                  this will return an empty list.
+        :param dataset_id: The dataset over which to execute the transaction.
         """
-        lookup_request = datastore_pb.LookupRequest()
-
-        single_key = isinstance(key_pbs, datastore_pb.Key)
 
-        if single_key:
-            key_pbs = [key_pbs]
-
-        for key_pb in key_pbs:
-            lookup_request.key.add().CopyFrom(key_pb)
+        if self.transaction():
+            raise ValueError('Cannot start a transaction with another already '
+                             'in progress.')
 
-        lookup_response = self._rpc(dataset_id, 'lookup', lookup_request,
-                                    datastore_pb.LookupResponse)
+        request = datastore_pb.BeginTransactionRequest()
 
-        results = [result.entity for result in lookup_response.found]
+        if serializable:
+            request.isolation_level = (
+                datastore_pb.BeginTransactionRequest.SERIALIZABLE)
+        else:
+            request.isolation_level = (
+                datastore_pb.BeginTransactionRequest.SNAPSHOT)
 
-        if single_key:
-            if results:
-                return results[0]
-            else:
-                return None
+        response = self._rpc(dataset_id, 'beginTransaction', request,
+                             datastore_pb.BeginTransactionResponse)
 
-        return results
+        return response.transaction
 
     def commit(self, dataset_id, mutation_pb):
         """Commit dataset mutations in context of current transation (if any).
 
+        Maps the ``DatastoreService.Commit`` protobuf RPC.
+
         :type dataset_id: string
         :param dataset_id: The dataset in which to perform the changes.
 
@@ -339,6 +329,48 @@ def commit(self, dataset_id, mutation_pb):
                              datastore_pb.CommitResponse)
         return response.mutation_result
 
+    def rollback(self, dataset_id):
+        """Rollback the connection's existing transaction.
+
+        Maps the ``DatastoreService.Rollback`` protobuf RPC.
+
+        Raises a ``ValueError``
+        if the connection isn't currently in a transaction.
+
+        :type dataset_id: string
+        :param dataset_id: The dataset to which the transaction belongs.
+        """
+        if not self.transaction() or not self.transaction().id():
+            raise ValueError('No transaction to rollback.')
+
+        request = datastore_pb.RollbackRequest()
+        request.transaction = self.transaction().id()
+        # Nothing to do with this response, so just execute the method.
+        self._rpc(dataset_id, 'rollback', request,
+                  datastore_pb.RollbackResponse)
+
+    def allocate_ids(self, dataset_id, key_pbs):
+        """Obtain backend-generated IDs for a set of keys.
+
+        Maps the ``DatastoreService.AllocateIds`` protobuf RPC.
+
+        :type dataset_id: string
+        :param dataset_id: The dataset to which the transaction belongs.
+
+        :type key_pbs: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
+        :param key_pbs: The keys for which the backend should allocate IDs.
+
+        :rtype: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
+        :returns: An equal number of keys,  with IDs filled in by the backend.
+        """
+        request = datastore_pb.AllocateIdsRequest()
+        for key_pb in key_pbs:
+            request.key.add().CopyFrom(key_pb)
+        # Nothing to do with this response, so just execute the method.
+        response = self._rpc(dataset_id, 'allocateIds', request,
+                             datastore_pb.AllocateIdsResponse)
+        return list(response.key)
+
     def save_entity(self, dataset_id, key_pb, properties):
         """Save an entity to the Cloud Datastore with the provided properties.
 
diff --git a/gcloud/datastore/transaction.py b/gcloud/datastore/transaction.py
--- a/gcloud/datastore/transaction.py
+++ b/gcloud/datastore/transaction.py
@@ -210,7 +210,7 @@ def rollback(self):
         - Sets the current connection's transaction reference to None.
         - Sets the current transaction's ID to None.
         """
-        self.connection().rollback_transaction(self.dataset().id())
+        self.connection().rollback(self.dataset().id())
         self.connection().transaction(None)
         self._id = None
 


Test patch:
diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py
--- a/gcloud/datastore/test_connection.py
+++ b/gcloud/datastore/test_connection.py
@@ -192,19 +192,13 @@ def test_dataset(self):
         self.assertTrue(dataset.connection() is conn)
         self.assertEqual(dataset.id(), DATASET_ID)
 
-    def test_begin_transaction_w_existing_transaction(self):
-        DATASET_ID = 'DATASET'
-        conn = self._makeOne()
-        conn.transaction(object())
-        self.assertRaises(ValueError, conn.begin_transaction, DATASET_ID)
-
-    def test_begin_transaction_default_serialize(self):
+    def test_lookup_single_key_empty_response(self):
         from gcloud.datastore.connection import datastore_pb
+        from gcloud.datastore.key import Key
 
         DATASET_ID = 'DATASET'
-        TRANSACTION = 'TRANSACTION'
-        rsp_pb = datastore_pb.BeginTransactionResponse()
-        rsp_pb.transaction = TRANSACTION
+        key_pb = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
+        rsp_pb = datastore_pb.LookupResponse()
         conn = self._makeOne()
         URI = '/'.join([
             conn.API_BASE_URL,
@@ -212,28 +206,33 @@ def test_begin_transaction_default_serialize(self):
             conn.API_VERSION,
             'datasets',
             DATASET_ID,
-            'beginTransaction',
+            'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.begin_transaction(DATASET_ID), TRANSACTION)
+        self.assertEqual(conn.lookup(DATASET_ID, key_pb), None)
         cw = http._called_with
         self.assertEqual(cw['uri'], URI)
         self.assertEqual(cw['method'], 'POST')
         self.assertEqual(cw['headers']['Content-Type'],
                          'application/x-protobuf')
         self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
-        rq_class = datastore_pb.BeginTransactionRequest
+        rq_class = datastore_pb.LookupRequest
         request = rq_class()
         request.ParseFromString(cw['body'])
-        self.assertEqual(request.isolation_level, rq_class.SNAPSHOT)
+        keys = list(request.key)
+        self.assertEqual(len(keys), 1)
+        self.assertEqual(keys[0], key_pb)
 
-    def test_begin_transaction_explicit_serialize(self):
+    def test_lookup_single_key_nonempty_response(self):
         from gcloud.datastore.connection import datastore_pb
+        from gcloud.datastore.key import Key
 
         DATASET_ID = 'DATASET'
-        TRANSACTION = 'TRANSACTION'
-        rsp_pb = datastore_pb.BeginTransactionResponse()
-        rsp_pb.transaction = TRANSACTION
+        key_pb = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
+        rsp_pb = datastore_pb.LookupResponse()
+        entity = datastore_pb.Entity()
+        entity.key.CopyFrom(key_pb)
+        rsp_pb.found.add(entity=entity)
         conn = self._makeOne()
         URI = '/'.join([
             conn.API_BASE_URL,
@@ -241,71 +240,57 @@ def test_begin_transaction_explicit_serialize(self):
             conn.API_VERSION,
             'datasets',
             DATASET_ID,
-            'beginTransaction',
+            'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.begin_transaction(DATASET_ID, True), TRANSACTION)
+        found = conn.lookup(DATASET_ID, key_pb)
+        self.assertEqual(found.key.path_element[0].kind, 'Kind')
+        self.assertEqual(found.key.path_element[0].id, 1234)
         cw = http._called_with
         self.assertEqual(cw['uri'], URI)
         self.assertEqual(cw['method'], 'POST')
         self.assertEqual(cw['headers']['Content-Type'],
                          'application/x-protobuf')
         self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
-        rq_class = datastore_pb.BeginTransactionRequest
+        rq_class = datastore_pb.LookupRequest
         request = rq_class()
         request.ParseFromString(cw['body'])
-        self.assertEqual(request.isolation_level, rq_class.SERIALIZABLE)
-
-    def test_rollback_transaction_wo_existing_transaction(self):
-        DATASET_ID = 'DATASET'
-        conn = self._makeOne()
-        self.assertRaises(ValueError,
-                          conn.rollback_transaction, DATASET_ID)
-
-    def test_rollback_transaction_w_existing_transaction_no_id(self):
-
-        class Xact(object):
-
-            def id(self):
-                return None
-        DATASET_ID = 'DATASET'
-        conn = self._makeOne()
-        conn.transaction(Xact())
-        self.assertRaises(ValueError,
-                          conn.rollback_transaction, DATASET_ID)
+        keys = list(request.key)
+        self.assertEqual(len(keys), 1)
+        self.assertEqual(keys[0], key_pb)
 
-    def test_rollback_transaction_ok(self):
+    def test_lookup_multiple_keys_empty_response(self):
         from gcloud.datastore.connection import datastore_pb
-        DATASET_ID = 'DATASET'
-        TRANSACTION = 'xact'
-
-        class Xact(object):
+        from gcloud.datastore.key import Key
 
-            def id(self):
-                return TRANSACTION
-        rsp_pb = datastore_pb.RollbackResponse()
+        DATASET_ID = 'DATASET'
+        key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
+        key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf()
+        rsp_pb = datastore_pb.LookupResponse()
         conn = self._makeOne()
-        conn.transaction(Xact())
         URI = '/'.join([
             conn.API_BASE_URL,
             'datastore',
             conn.API_VERSION,
             'datasets',
             DATASET_ID,
-            'rollback',
+            'lookup',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.rollback_transaction(DATASET_ID), None)
+        self.assertEqual(conn.lookup(DATASET_ID, [key_pb1, key_pb2]), [])
         cw = http._called_with
         self.assertEqual(cw['uri'], URI)
         self.assertEqual(cw['method'], 'POST')
         self.assertEqual(cw['headers']['Content-Type'],
                          'application/x-protobuf')
         self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
-        rq_class = datastore_pb.RollbackRequest
+        rq_class = datastore_pb.LookupRequest
         request = rq_class()
         request.ParseFromString(cw['body'])
-        self.assertEqual(request.transaction, TRANSACTION)
+        keys = list(request.key)
+        self.assertEqual(len(keys), 2)
+        self.assertEqual(keys[0], key_pb1)
+        self.assertEqual(keys[1], key_pb2)
 
     def test_run_query_wo_namespace_empty_result(self):
         from gcloud.datastore.connection import datastore_pb
@@ -378,47 +363,19 @@ def test_run_query_w_namespace_nonempty_result(self):
         self.assertEqual(request.partition_id.namespace, 'NS')
         self.assertEqual(request.query, q_pb)
 
-    def test_lookup_single_key_empty_response(self):
-        from gcloud.datastore.connection import datastore_pb
-        from gcloud.datastore.key import Key
-
+    def test_begin_transaction_w_existing_transaction(self):
         DATASET_ID = 'DATASET'
-        key_pb = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
-        rsp_pb = datastore_pb.LookupResponse()
         conn = self._makeOne()
-        URI = '/'.join([
-            conn.API_BASE_URL,
-            'datastore',
-            conn.API_VERSION,
-            'datasets',
-            DATASET_ID,
-            'lookup',
-        ])
-        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.lookup(DATASET_ID, key_pb), None)
-        cw = http._called_with
-        self.assertEqual(cw['uri'], URI)
-        self.assertEqual(cw['method'], 'POST')
-        self.assertEqual(cw['headers']['Content-Type'],
-                         'application/x-protobuf')
-        self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
-        rq_class = datastore_pb.LookupRequest
-        request = rq_class()
-        request.ParseFromString(cw['body'])
-        keys = list(request.key)
-        self.assertEqual(len(keys), 1)
-        self.assertEqual(keys[0], key_pb)
+        conn.transaction(object())
+        self.assertRaises(ValueError, conn.begin_transaction, DATASET_ID)
 
-    def test_lookup_single_key_nonempty_response(self):
+    def test_begin_transaction_default_serialize(self):
         from gcloud.datastore.connection import datastore_pb
-        from gcloud.datastore.key import Key
 
         DATASET_ID = 'DATASET'
-        key_pb = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
-        rsp_pb = datastore_pb.LookupResponse()
-        entity = datastore_pb.Entity()
-        entity.key.CopyFrom(key_pb)
-        rsp_pb.found.add(entity=entity)
+        TRANSACTION = 'TRANSACTION'
+        rsp_pb = datastore_pb.BeginTransactionResponse()
+        rsp_pb.transaction = TRANSACTION
         conn = self._makeOne()
         URI = '/'.join([
             conn.API_BASE_URL,
@@ -426,33 +383,28 @@ def test_lookup_single_key_nonempty_response(self):
             conn.API_VERSION,
             'datasets',
             DATASET_ID,
-            'lookup',
+            'beginTransaction',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        found = conn.lookup(DATASET_ID, key_pb)
-        self.assertEqual(found.key.path_element[0].kind, 'Kind')
-        self.assertEqual(found.key.path_element[0].id, 1234)
+        self.assertEqual(conn.begin_transaction(DATASET_ID), TRANSACTION)
         cw = http._called_with
         self.assertEqual(cw['uri'], URI)
         self.assertEqual(cw['method'], 'POST')
         self.assertEqual(cw['headers']['Content-Type'],
                          'application/x-protobuf')
         self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
-        rq_class = datastore_pb.LookupRequest
+        rq_class = datastore_pb.BeginTransactionRequest
         request = rq_class()
         request.ParseFromString(cw['body'])
-        keys = list(request.key)
-        self.assertEqual(len(keys), 1)
-        self.assertEqual(keys[0], key_pb)
+        self.assertEqual(request.isolation_level, rq_class.SNAPSHOT)
 
-    def test_lookup_multiple_keys_empty_response(self):
+    def test_begin_transaction_explicit_serialize(self):
         from gcloud.datastore.connection import datastore_pb
-        from gcloud.datastore.key import Key
 
         DATASET_ID = 'DATASET'
-        key_pb1 = Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
-        key_pb2 = Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf()
-        rsp_pb = datastore_pb.LookupResponse()
+        TRANSACTION = 'TRANSACTION'
+        rsp_pb = datastore_pb.BeginTransactionResponse()
+        rsp_pb.transaction = TRANSACTION
         conn = self._makeOne()
         URI = '/'.join([
             conn.API_BASE_URL,
@@ -460,23 +412,20 @@ def test_lookup_multiple_keys_empty_response(self):
             conn.API_VERSION,
             'datasets',
             DATASET_ID,
-            'lookup',
+            'beginTransaction',
         ])
         http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
-        self.assertEqual(conn.lookup(DATASET_ID, [key_pb1, key_pb2]), [])
+        self.assertEqual(conn.begin_transaction(DATASET_ID, True), TRANSACTION)
         cw = http._called_with
         self.assertEqual(cw['uri'], URI)
         self.assertEqual(cw['method'], 'POST')
         self.assertEqual(cw['headers']['Content-Type'],
                          'application/x-protobuf')
         self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
-        rq_class = datastore_pb.LookupRequest
+        rq_class = datastore_pb.BeginTransactionRequest
         request = rq_class()
         request.ParseFromString(cw['body'])
-        keys = list(request.key)
-        self.assertEqual(len(keys), 2)
-        self.assertEqual(keys[0], key_pb1)
-        self.assertEqual(keys[1], key_pb2)
+        self.assertEqual(request.isolation_level, rq_class.SERIALIZABLE)
 
     def test_commit_wo_transaction(self):
         from gcloud.datastore.connection import datastore_pb
@@ -560,6 +509,123 @@ def id(self):
         self.assertEqual(request.mutation, mutation)
         self.assertEqual(request.mode, rq_class.TRANSACTIONAL)
 
+    def test_rollback_wo_existing_transaction(self):
+        DATASET_ID = 'DATASET'
+        conn = self._makeOne()
+        self.assertRaises(ValueError,
+                          conn.rollback, DATASET_ID)
+
+    def test_rollback_w_existing_transaction_no_id(self):
+
+        class Xact(object):
+
+            def id(self):
+                return None
+        DATASET_ID = 'DATASET'
+        conn = self._makeOne()
+        conn.transaction(Xact())
+        self.assertRaises(ValueError,
+                          conn.rollback, DATASET_ID)
+
+    def test_rollback_ok(self):
+        from gcloud.datastore.connection import datastore_pb
+        DATASET_ID = 'DATASET'
+        TRANSACTION = 'xact'
+
+        class Xact(object):
+
+            def id(self):
+                return TRANSACTION
+        rsp_pb = datastore_pb.RollbackResponse()
+        conn = self._makeOne()
+        conn.transaction(Xact())
+        URI = '/'.join([
+            conn.API_BASE_URL,
+            'datastore',
+            conn.API_VERSION,
+            'datasets',
+            DATASET_ID,
+            'rollback',
+        ])
+        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
+        self.assertEqual(conn.rollback(DATASET_ID), None)
+        cw = http._called_with
+        self.assertEqual(cw['uri'], URI)
+        self.assertEqual(cw['method'], 'POST')
+        self.assertEqual(cw['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
+        rq_class = datastore_pb.RollbackRequest
+        request = rq_class()
+        request.ParseFromString(cw['body'])
+        self.assertEqual(request.transaction, TRANSACTION)
+
+    def test_allocate_ids_empty(self):
+        from gcloud.datastore.connection import datastore_pb
+
+        DATASET_ID = 'DATASET'
+        rsp_pb = datastore_pb.AllocateIdsResponse()
+        conn = self._makeOne()
+        URI = '/'.join([
+            conn.API_BASE_URL,
+            'datastore',
+            conn.API_VERSION,
+            'datasets',
+            DATASET_ID,
+            'allocateIds',
+        ])
+        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
+        self.assertEqual(conn.allocate_ids(DATASET_ID, []), [])
+        cw = http._called_with
+        self.assertEqual(cw['uri'], URI)
+        self.assertEqual(cw['method'], 'POST')
+        self.assertEqual(cw['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
+        rq_class = datastore_pb.AllocateIdsRequest
+        request = rq_class()
+        request.ParseFromString(cw['body'])
+        self.assertEqual(list(request.key), [])
+
+    def test_allocate_ids_non_empty(self):
+        from gcloud.datastore.connection import datastore_pb
+        from gcloud.datastore.key import Key
+
+        DATASET_ID = 'DATASET'
+        before_key_pbs = [
+            Key(path=[{'kind': 'Kind'}]).to_protobuf(),
+            Key(path=[{'kind': 'Kind'}]).to_protobuf(),
+            ]
+        after_key_pbs = [
+            Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf(),
+            Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf(),
+            ]
+        rsp_pb = datastore_pb.AllocateIdsResponse()
+        rsp_pb.key.add().CopyFrom(after_key_pbs[0])
+        rsp_pb.key.add().CopyFrom(after_key_pbs[1])
+        conn = self._makeOne()
+        URI = '/'.join([
+            conn.API_BASE_URL,
+            'datastore',
+            conn.API_VERSION,
+            'datasets',
+            DATASET_ID,
+            'allocateIds',
+        ])
+        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
+        self.assertEqual(conn.allocate_ids(DATASET_ID, before_key_pbs),
+                         after_key_pbs)
+        cw = http._called_with
+        self.assertEqual(cw['uri'], URI)
+        self.assertEqual(cw['method'], 'POST')
+        self.assertEqual(cw['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
+        rq_class = datastore_pb.AllocateIdsRequest
+        request = rq_class()
+        request.ParseFromString(cw['body'])
+        self.assertEqual(list(request.key), before_key_pbs)
+
     def test_save_entity_wo_transaction_w_upsert(self):
         from gcloud.datastore.connection import datastore_pb
         from gcloud.datastore.key import Key
diff --git a/gcloud/datastore/test_transaction.py b/gcloud/datastore/test_transaction.py
--- a/gcloud/datastore/test_transaction.py
+++ b/gcloud/datastore/test_transaction.py
@@ -174,7 +174,7 @@ def begin_transaction(self, dataset_id):
         self._begun = dataset_id
         return self._xact_id
 
-    def rollback_transaction(self, dataset_id):
+    def rollback(self, dataset_id):
         self._rolled_back = dataset_id
 
     def commit(self, dataset_id, mutation):