Making datastore Connection.begin_transaction() return low-level protobuf.

Author: dhermes

datastore/google/cloud/datastore/_http.py

@@ -399,12 +399,11 @@ def begin_transaction(self, project):
         :type project: str
         :param project: The project to which the transaction applies.
 
-        :rtype: bytes
+        :rtype: :class:`.datastore_pb2.BeginTransactionResponse`
         :returns: The serialized transaction that was begun.
         """
         request = _datastore_pb2.BeginTransactionRequest()
-        response = self._datastore_api.begin_transaction(project, request)
-        return response.transaction
+        return self._datastore_api.begin_transaction(project, request)
 
     def commit(self, project, request, transaction_id):
         """Commit mutations in context of current transaction (if any).

datastore/google/cloud/datastore/transaction.py

@@ -184,8 +184,9 @@ def begin(self):
         """
         super(Transaction, self).begin()
         try:
-            self._id = self._client._connection.begin_transaction(
+            response_pb = self._client._connection.begin_transaction(
                 self.project)
+            self._id = response_pb.transaction
         except:  # noqa: E722 do not use bare except, specify exception instead
             self._status = self._ABORTED
             raise

datastore/unit_tests/test__http.py

@@ -672,25 +672,33 @@ def test_run_query_w_namespace_nonempty_result(self):
     def test_begin_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
 
-        PROJECT = 'PROJECT'
-        TRANSACTION = b'TRANSACTION'
+        project = 'PROJECT'
+        transaction = b'TRANSACTION'
         rsp_pb = datastore_pb2.BeginTransactionResponse()
-        rsp_pb.transaction = TRANSACTION
+        rsp_pb.transaction = transaction
+
+        # Create mock HTTP and client with response.
         http = Http({'status': '200'}, rsp_pb.SerializeToString())
         client = mock.Mock(_http=http, spec=['_http'])
+
+        # Make request.
         conn = self._make_one(client)
-        URI = '/'.join([
+        response = conn.begin_transaction(project)
+
+        # Check the result and verify the callers.
+        self.assertEqual(response, rsp_pb)
+        uri = '/'.join([
             conn.api_base_url,
             conn.API_VERSION,
             'projects',
-            PROJECT + ':beginTransaction',
+            project + ':beginTransaction',
         ])
-        self.assertEqual(conn.begin_transaction(PROJECT), TRANSACTION)
         cw = http._called_with
-        self._verify_protobuf_call(cw, URI, conn)
-        rq_class = datastore_pb2.BeginTransactionRequest
-        request = rq_class()
+        self._verify_protobuf_call(cw, uri, conn)
+        request = datastore_pb2.BeginTransactionRequest()
         request.ParseFromString(cw['body'])
+        # The RPC-over-HTTP request does not set the project in the request.
+        self.assertEqual(request.project_id, u'')
 
     def test_commit_wo_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2

datastore/unit_tests/test_transaction.py

@@ -206,9 +206,12 @@ def __init__(self, xact_id=123, keys=()):
             mutation_results=mutation_results)
 
     def begin_transaction(self, project):
+        import mock
+
         self._begun = project
         if self._side_effect is None:
-            return self._xact_id
+            return mock.Mock(
+                transaction=self._xact_id, spec=['transaction'])
         else:
             raise self._side_effect