fix: changes for system tests

google/cloud/datastore/_gapic.py

@@ -19,8 +19,8 @@
 
 from google.cloud._helpers import make_secure_channel
 from google.cloud._http import DEFAULT_USER_AGENT
-from google.cloud.datastore_v1.gapic import datastore_client
-
+from google.cloud.datastore_v1.services.datastore import client as datastore_client
+from google.cloud.datastore_v1.services.datastore.transports import grpc
 
 def make_datastore_api(client):
     """Create an instance of the GAPIC Datastore API.
@@ -38,6 +38,49 @@ def make_datastore_api(client):
     else:
         channel = insecure_channel(host)
 
+    #from .base import DatastoreTransport
+
+    transport = grpc.DatastoreGrpcTransport(channel=channel)
+
     return datastore_client.DatastoreClient(
-        channel=channel, client_info=client._client_info
+        transport=transport, client_info=client._client_info
     )
+
+    """
+    Instantiate the datastore client.
+
+        Args:
+            credentials (Optional[google.auth.credentials.Credentials]): The
+                authorization credentials to attach to requests. These
+                credentials identify the application to the service; if none
+                are specified, the client will attempt to ascertain the
+                credentials from the environment.
+            transport (Union[str, ~.DatastoreTransport]): The
+                transport to use. If set to None, a transport is chosen
+                automatically.
+            client_options (ClientOptions): Custom options for the client. It
+                won't take effect if a ``transport`` instance is provided.
+                (1) The ``api_endpoint`` property can be used to override the
+                default endpoint provided by the client. GOOGLE_API_USE_MTLS_ENDPOINT
+                environment variable can also be used to override the endpoint:
+                "always" (always use the default mTLS endpoint), "never" (always
+                use the default regular endpoint) and "auto" (auto switch to the
+                default mTLS endpoint if client certificate is present, this is
+                the default value). However, the ``api_endpoint`` property takes
+                precedence if provided.
+                (2) If GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable
+                is "true", then the ``client_cert_source`` property can be used
+                to provide client certificate for mutual TLS transport. If
+                not provided, the default SSL client certificate will be used if
+                present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not
+                set, no client certificate will be used.
+            client_info (google.api_core.gapic_v1.client_info.ClientInfo):	
+                The client info used to send a user-agent string along with	
+                API requests. If ``None``, then default info will be used.	
+                Generally, you only need to set this if you're developing	
+                your own client library.
+
+        Raises:
+            google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport
+                creation failed for any reason.
+                """

google/cloud/datastore/_http.py

@@ -110,8 +110,9 @@ def _rpc(http, project, method, base_url, client_info, request_pb, response_pb_c
     """
     req_data = request_pb.SerializeToString()
     response = _request(http, project, method, req_data, base_url, client_info)
-    return response_pb_cls._meta._pb.FromString(response)
-
+    from_string =  response_pb_cls._meta._pb.FromString(response)
+    deserialized = response_pb_cls.deserialize(response)
+    return deserialized
 
 def build_api_url(project, method, base_url):
     """Construct the URL for a particular API call.

google/cloud/datastore/batch.py

@@ -255,8 +255,15 @@ def _commit(self, retry, timeout):
             kwargs["timeout"] = timeout
 
         commit_response_pb = self._client._datastore_api.commit(
-            self.project, mode, self._mutations, transaction=self._id, **kwargs
-        )
+            request = {
+               'project_id': self.project,
+               'mode': mode,
+               'transaction': self._id,
+               'mutations': self._mutations,
+               **kwargs
+            },
+
+        )._pb
         _, updated_keys = _parse_commit_response(commit_response_pb)
         # If the back-end returns without error, we are guaranteed that
         # ``commit`` will return keys that match (length and

google/cloud/datastore/client.py

@@ -185,7 +185,12 @@ def _extended_lookup(
     while loop_num < _MAX_LOOPS:  # loop against possible deferred.
         loop_num += 1
         lookup_response = datastore_api.lookup(
-            project, key_pbs, read_options=read_options, **kwargs
+            request = {
+                "project_id": project,
+                "keys": key_pbs,
+                "read_options": read_options,
+                **kwargs
+            },
         )
 
         # Accumulate the new results.
@@ -535,7 +540,7 @@ def get_multi(
                 helpers.key_from_protobuf(deferred_pb) for deferred_pb in deferred
             ]
 
-        return [helpers.entity_from_protobuf(entity_pb) for entity_pb in entity_pbs]
+        return [helpers.entity_from_protobuf(entity_pb._pb) for entity_pb in entity_pbs]
 
     def put(self, entity, retry=None, timeout=None):
         """Save an entity in the Cloud Datastore.
@@ -702,7 +707,11 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
 
         response_pb = self._datastore_api.allocate_ids(
-            incomplete_key.project, incomplete_key_pbs, **kwargs
+            request = {
+                "project_id": incomplete_key.project,
+                "keys": incomplete_key_pbs,
+                **kwargs
+            },
         )
         allocated_ids = [
             allocated_key_pb.path[-1].id for allocated_key_pb in response_pb.keys
@@ -871,8 +880,13 @@ def reserve_ids_sequential(self, complete_key, num_ids, retry=None, timeout=None
             key_pbs.append(key.to_protobuf())
 
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
-        self._datastore_api.reserve_ids(complete_key.project, key_pbs, **kwargs)
-
+        self._datastore_api.reserve_ids(
+            request={
+                "project_id": complete_key.project,
+                "keys": key_pbs,
+                **kwargs
+            }
+        )
         return None
 
     def reserve_ids(self, complete_key, num_ids, retry=None, timeout=None):
@@ -921,6 +935,12 @@ def reserve_ids_multi(self, complete_keys, retry=None, timeout=None):
 
         kwargs = _make_retry_timeout_kwargs(retry, timeout)
         key_pbs = [key.to_protobuf() for key in complete_keys]
-        self._datastore_api.reserve_ids(complete_keys[0].project, key_pbs, **kwargs)
+        self._datastore_api.reserve_ids(
+            request = {
+                "project_id": complete_keys[0].project,
+                "keys": key_pbs,
+                **kwargs
+            }
+        )
 
         return None

google/cloud/datastore/query.py

@@ -575,22 +575,14 @@ def _next_page(self):
         if self._timeout is not None:
             kwargs["timeout"] = self._timeout
 
-        # response_pb = self.client._datastore_api.run_query(
-        #     request={
-        #         "project_id": self._query.project,
-        #         "partition_id": partition_id,
-        #         "read_options": read_options,
-        #         "query": query_pb,
-        #         # "gql_query": query_pb,
-        #     },
-        #     **kwargs
-        # )
         response_pb = self.client._datastore_api.run_query(
-            self._query.project,
-            partition_id,
-            read_options,
-            query=query_pb,
-            **kwargs
+            request = {
+                "project_id": self._query.project,
+                "partition_id": partition_id,
+                "read_options": read_options,
+                "query": query_pb,
+                **kwargs
+            }
         )
 
         while (
@@ -603,21 +595,14 @@ def _next_page(self):
             # more than 1000 skipped results in a query.
             query_pb.start_cursor = response_pb.batch.skipped_cursor
             query_pb.offset -= response_pb.batch.skipped_results
-            # response_pb = self.client._datastore_api.run_query(
-            #     request={
-            #         "project_id": self._query.project,
-            #         "partition_id": partition_id,
-            #         "read_options": read_options,
-            #         "query": kwargs,
-            #         "gql_query": query_pb,
-            #     }
-            # )
             response_pb = self.client._datastore_api.run_query(
-                self._query.project,
-                partition_id,
-                read_options,
-                query=query_pb,
-                **kwargs,
+                request={
+                    "project_id": self._query.project,
+                    "partition_id": partition_id,
+                    "read_options": read_options,
+                    "query": query_pb,
+                    **kwargs
+                }
             )
         entity_pbs = self._process_query_results(response_pb)
         return page_iterator.Page(self, entity_pbs, self.item_to_value)
@@ -702,7 +687,7 @@ def _item_to_entity(iterator, entity_pb):
     :rtype: :class:`~google.cloud.datastore.entity.Entity`
     :returns: The next entity in the page.
     """
-    return helpers.entity_from_protobuf(entity_pb)
+    return helpers.entity_from_protobuf(entity_pb._pb)
 
 
 # pylint: enable=unused-argument

google/cloud/datastore/transaction.py

@@ -233,7 +233,7 @@ def begin(self, retry=None, timeout=None):
 
         try:
             response_pb = self._client._datastore_api.begin_transaction(
-                self.project, **kwargs
+                project_id = self.project, **kwargs
             )
             self._id = response_pb.transaction
         except:  # noqa: E722 do not use bare except, specify exception instead

tests/unit/test_batch.py

@@ -525,7 +525,7 @@ def _make_commit_response(*new_key_ids):
     from google.cloud.datastore_v1.types import datastore as datastore_pb2
 
     mutation_results = [_make_mutation(key_id) for key_id in new_key_ids]
-    return datastore_pb2.CommitResponse(mutation_results=mutation_results)
+    return datastore_pb2.CommitResponse(mutation_results=mutation_results)._pb
 
 
 def _make_datastore_api(*new_key_ids):

tests/unit/test_query.py

@@ -579,8 +579,8 @@ def _next_page_helper(self, txn_id=None, retry=None, timeout=None):
                 "partition_id": partition_id,
                 "read_options": read_options,
                 "query": empty_query,
+                **kwargs
             },
-            **kwargs
         )
 
     def test__next_page(self):