Merge pull request #3134 from dhermes/ds-move-begin_transaction-to-GAPIC

Using GAPIC datastore object (and an HTTP equivalent) for begin_transaction.

Author: dhermes

datastore/google/cloud/datastore/_gax.py

@@ -152,24 +152,6 @@ def run_query(self, project, request_pb):
         with _grpc_catch_rendezvous():
             return self._stub.RunQuery(request_pb)
 
-    def begin_transaction(self, project, request_pb):
-        """Perform a ``beginTransaction`` request.
-
-        :type project: str
-        :param project: The project to connect to. This is
-                        usually your project name in the cloud console.
-
-        :type request_pb:
-            :class:`.datastore_pb2.BeginTransactionRequest`
-        :param request_pb: The request protobuf object.
-
-        :rtype: :class:`.datastore_pb2.BeginTransactionResponse`
-        :returns: The returned protobuf response object.
-        """
-        request_pb.project_id = project
-        with _grpc_catch_rendezvous():
-            return self._stub.BeginTransaction(request_pb)
-
 
 class GAPICDatastoreAPI(datastore_client.DatastoreClient):
     """An API object that sends proto-over-gRPC requests.

datastore/google/cloud/datastore/_http.py

@@ -198,24 +198,6 @@ def run_query(self, project, request_pb):
                     self.connection.api_base_url,
                     request_pb, _datastore_pb2.RunQueryResponse)
 
-    def begin_transaction(self, project, request_pb):
-        """Perform a ``beginTransaction`` request.
-
-        :type project: str
-        :param project: The project to connect to. This is
-                        usually your project name in the cloud console.
-
-        :type request_pb:
-            :class:`.datastore_pb2.BeginTransactionRequest`
-        :param request_pb: The request protobuf object.
-
-        :rtype: :class:`.datastore_pb2.BeginTransactionResponse`
-        :returns: The returned protobuf response object.
-        """
-        return _rpc(self.connection.http, project, 'beginTransaction',
-                    self.connection.api_base_url,
-                    request_pb, _datastore_pb2.BeginTransactionResponse)
-
 
 class Connection(connection_module.Connection):
     """A connection to the Google Cloud Datastore via the Protobuf API.
@@ -335,20 +317,6 @@ def run_query(self, project, query_pb, namespace=None,
         request.query.CopyFrom(query_pb)
         return self._datastore_api.run_query(project, request)
 
-    def begin_transaction(self, project):
-        """Begin a transaction.
-
-        Maps the ``DatastoreService.BeginTransaction`` protobuf RPC.
-
-        :type project: str
-        :param project: The project to which the transaction applies.
-
-        :rtype: :class:`.datastore_pb2.BeginTransactionResponse`
-        :returns: The serialized transaction that was begun.
-        """
-        request = _datastore_pb2.BeginTransactionRequest()
-        return self._datastore_api.begin_transaction(project, request)
-
 
 class HTTPDatastoreAPI(object):
     """An API object that sends proto-over-HTTP requests.
@@ -362,6 +330,21 @@ class HTTPDatastoreAPI(object):
     def __init__(self, client):
         self.client = client
 
+    def begin_transaction(self, project):
+        """Perform a ``beginTransaction`` request.
+
+        :type project: str
+        :param project: The project to connect to. This is
+                        usually your project name in the cloud console.
+
+        :rtype: :class:`.datastore_pb2.BeginTransactionResponse`
+        :returns: The returned protobuf response object.
+        """
+        request_pb = _datastore_pb2.BeginTransactionRequest()
+        return _rpc(self.client._http, project, 'beginTransaction',
+                    self.client._base_url,
+                    request_pb, _datastore_pb2.BeginTransactionResponse)
+
     def commit(self, project, mode, mutations, transaction=None):
         """Perform a ``commit`` request.
 

datastore/google/cloud/datastore/transaction.py

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

datastore/unit_tests/test__gax.py

@@ -240,20 +240,6 @@ def test_run_query_invalid_argument(self):
         exc = GrpcRendezvous(exc_state, None, None, None)
         self._run_query_failure_helper(exc, BadRequest)
 
-    def test_begin_transaction(self):
-        return_val = object()
-        stub = _GRPCStub(return_val)
-        datastore_api, _ = self._make_one(stub=stub)
-
-        request_pb = mock.Mock(project_id=None, spec=['project_id'])
-        project = 'PROJECT'
-        result = datastore_api.begin_transaction(project, request_pb)
-        self.assertIs(result, return_val)
-        self.assertEqual(request_pb.project_id, project)
-        self.assertEqual(
-            stub.method_calls,
-            [(request_pb, 'BeginTransaction')])
-
 
 @unittest.skipUnless(_HAVE_GRPC, 'No gRPC')
 class TestGAPICDatastoreAPI(unittest.TestCase):
@@ -338,6 +324,3 @@ def Lookup(self, request_pb):
 
     def RunQuery(self, request_pb):
         return self._method(request_pb, 'RunQuery')
-
-    def BeginTransaction(self, request_pb):
-        return self._method(request_pb, 'BeginTransaction')

datastore/unit_tests/test__http.py

@@ -634,6 +634,23 @@ def test_run_query_w_namespace_nonempty_result(self):
         self.assertEqual(request.partition_id.namespace_id, namespace)
         self.assertEqual(request.query, q_pb)
 
+
+class TestHTTPDatastoreAPI(unittest.TestCase):
+
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.datastore._http import HTTPDatastoreAPI
+
+        return HTTPDatastoreAPI
+
+    def _make_one(self, *args, **kwargs):
+        return self._get_target_class()(*args, **kwargs)
+
+    def test_constructor(self):
+        client = object()
+        ds_api = self._make_one(client)
+        self.assertIs(ds_api.client, client)
+
     def test_begin_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
 
@@ -648,37 +665,20 @@ def test_begin_transaction(self):
             _http=http, _base_url='test.invalid', spec=['_http', '_base_url'])
 
         # Make request.
-        conn = self._make_one(client)
-        response = conn.begin_transaction(project)
+        ds_api = self._make_one(client)
+        response = ds_api.begin_transaction(project)
 
         # Check the result and verify the callers.
         self.assertEqual(response, rsp_pb)
         uri = _build_expected_url(
-            conn.api_base_url, project, 'beginTransaction')
+            client._base_url, project, 'beginTransaction')
         cw = http._called_with
         _verify_protobuf_call(self, cw, uri)
         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'')
 
-
-class TestHTTPDatastoreAPI(unittest.TestCase):
-
-    @staticmethod
-    def _get_target_class():
-        from google.cloud.datastore._http import HTTPDatastoreAPI
-
-        return HTTPDatastoreAPI
-
-    def _make_one(self, *args, **kwargs):
-        return self._get_target_class()(*args, **kwargs)
-
-    def test_constructor(self):
-        client = object()
-        ds_api = self._make_one(client)
-        self.assertIs(ds_api.client, client)
-
     def test_commit_wo_transaction(self):
         from google.cloud.proto.datastore.v1 import datastore_pb2
         from google.cloud.datastore.helpers import _new_value_pb