Datastore: Add client_info support to client.

datastore/google/cloud/datastore/_gapic.py

@@ -14,15 +14,12 @@
 
 """Helpers for making API requests via gapic / gRPC."""
 
-from google.api_core.gapic_v1 import client_info
-from google.cloud.datastore_v1.gapic import datastore_client
 from grpc import insecure_channel
 import six
 
 from google.cloud._helpers import make_secure_channel
 from google.cloud._http import DEFAULT_USER_AGENT
-
-from google.cloud.datastore import __version__
+from google.cloud.datastore_v1.gapic import datastore_client
 
 
 def make_datastore_api(client):
@@ -42,8 +39,5 @@ def make_datastore_api(client):
         channel = insecure_channel(host)
 
     return datastore_client.DatastoreClient(
-        channel=channel,
-        client_info=client_info.ClientInfo(
-            client_library_version=__version__, gapic_version=__version__
-        ),
+        channel=channel, client_info=client._client_info
     )

datastore/google/cloud/datastore/_http.py

@@ -20,8 +20,6 @@
 from google.cloud import exceptions
 from google.cloud.datastore_v1.proto import datastore_pb2 as _datastore_pb2
 
-from google.cloud.datastore import __version__
-
 
 DATASTORE_API_HOST = "datastore.googleapis.com"
 """Datastore API request host."""
@@ -32,10 +30,8 @@
 API_URL_TEMPLATE = "{api_base}/{api_version}/projects" "/{project}:{method}"
 """A template for the URL of a particular API call."""
 
-_CLIENT_INFO = connection_module.CLIENT_INFO_TEMPLATE.format(__version__)
-
 
-def _request(http, project, method, data, base_url):
+def _request(http, project, method, data, base_url, client_info):
     """Make a request over the Http transport to the Cloud Datastore API.
 
     :type http: :class:`requests.Session`
@@ -55,15 +51,19 @@ def _request(http, project, method, data, base_url):
     :type base_url: str
     :param base_url: The base URL where the API lives.
 
+    :type client_info: :class:`google.api_core.client_info.ClientInfo`
+    :param client_info: used to generate user agent.
+
     :rtype: str
     :returns: The string response content from the API call.
     :raises: :class:`google.cloud.exceptions.GoogleCloudError` if the
              response code is not 200 OK.
     """
+    user_agent = client_info.to_user_agent()
     headers = {
         "Content-Type": "application/x-protobuf",
-        "User-Agent": connection_module.DEFAULT_USER_AGENT,
-        connection_module.CLIENT_INFO_HEADER: _CLIENT_INFO,
+        "User-Agent": user_agent,
+        connection_module.CLIENT_INFO_HEADER: user_agent,
     }
     api_url = build_api_url(project, method, base_url)
 
@@ -78,7 +78,7 @@ def _request(http, project, method, data, base_url):
     return response.content
 
 
-def _rpc(http, project, method, base_url, request_pb, response_pb_cls):
+def _rpc(http, project, method, base_url, client_info, request_pb, response_pb_cls):
     """Make a protobuf RPC request.
 
     :type http: :class:`requests.Session`
@@ -94,6 +94,9 @@ def _rpc(http, project, method, base_url, request_pb, response_pb_cls):
     :type base_url: str
     :param base_url: The base URL where the API lives.
 
+    :type client_info: :class:`google.api_core.client_info.ClientInfo`
+    :param client_info: used to generate user agent.
+
     :type request_pb: :class:`google.protobuf.message.Message` instance
     :param request_pb: the protobuf instance representing the request.
 
@@ -106,7 +109,7 @@ def _rpc(http, project, method, base_url, request_pb, response_pb_cls):
     :returns: The RPC message parsed from the response.
     """
     req_data = request_pb.SerializeToString()
-    response = _request(http, project, method, req_data, base_url)
+    response = _request(http, project, method, req_data, base_url, client_info)
     return response_pb_cls.FromString(response)
 
 
@@ -172,6 +175,7 @@ def lookup(self, project_id, keys, read_options=None):
             project_id,
             "lookup",
             self.client._base_url,
+            self.client._client_info,
             request_pb,
             _datastore_pb2.LookupResponse,
         )
@@ -217,6 +221,7 @@ def run_query(
             project_id,
             "runQuery",
             self.client._base_url,
+            self.client._client_info,
             request_pb,
             _datastore_pb2.RunQueryResponse,
         )
@@ -240,6 +245,7 @@ def begin_transaction(self, project_id, transaction_options=None):
             project_id,
             "beginTransaction",
             self.client._base_url,
+            self.client._client_info,
             request_pb,
             _datastore_pb2.BeginTransactionResponse,
         )
@@ -278,6 +284,7 @@ def commit(self, project_id, mode, mutations, transaction=None):
             project_id,
             "commit",
             self.client._base_url,
+            self.client._client_info,
             request_pb,
             _datastore_pb2.CommitResponse,
         )
@@ -304,6 +311,7 @@ def rollback(self, project_id, transaction):
             project_id,
             "rollback",
             self.client._base_url,
+            self.client._client_info,
             request_pb,
             _datastore_pb2.RollbackResponse,
         )
@@ -327,6 +335,7 @@ def allocate_ids(self, project_id, keys):
             project_id,
             "allocateIds",
             self.client._base_url,
+            self.client._client_info,
             request_pb,
             _datastore_pb2.AllocateIdsResponse,
         )

datastore/google/cloud/datastore/client.py

@@ -18,6 +18,7 @@
 from google.cloud._helpers import _LocalStack
 from google.cloud._helpers import _determine_default_project as _base_default_project
 from google.cloud.client import ClientWithProject
+from google.cloud.datastore import __version__
 from google.cloud.datastore import helpers
 from google.cloud.datastore._http import HTTPDatastoreAPI
 from google.cloud.datastore.batch import Batch
@@ -32,17 +33,27 @@
 try:
     from google.cloud.datastore._gapic import make_datastore_api
 
-    _HAVE_GRPC = True
 except ImportError:  # pragma: NO COVER
+    from google.api_core import client_info
+
     make_datastore_api = None
     _HAVE_GRPC = False
+    _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__)
+else:
+    from google.api_core.gapic_v1 import client_info
+
+    _HAVE_GRPC = True
+    _CLIENT_INFO = client_info.ClientInfo(
+        client_library_version=__version__, gapic_version=__version__
+    )
 
 
 _MAX_LOOPS = 128
 """Maximum number of iterations to wait for deferred keys."""
 _DATASTORE_BASE_URL = "https://datastore.googleapis.com"
 """Datastore API request URL base."""
 
+
 _USE_GRPC = _HAVE_GRPC and not os.getenv(DISABLE_GRPC, False)
 
 
@@ -182,6 +193,14 @@ class Client(ClientWithProject):
                         passed), falls back to the default inferred from the
                         environment.
 
+    :type client_info: :class:`google.api_core.gapic_v1.client_info.ClientInfo`
+                       or :class:`google.api_core.client_info.ClientInfo`
+    :param client_info: (Optional) 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 library or partner tool.
+
     :type _http: :class:`~requests.Session`
     :param _http: (Optional) HTTP object to make requests. Can be any object
                   that defines ``request()`` with the same interface as
@@ -204,12 +223,19 @@ class Client(ClientWithProject):
     """The scopes required for authenticating as a Cloud Datastore consumer."""
 
     def __init__(
-        self, project=None, namespace=None, credentials=None, _http=None, _use_grpc=None
+        self,
+        project=None,
+        namespace=None,
+        credentials=None,
+        client_info=_CLIENT_INFO,
+        _http=None,
+        _use_grpc=None,
     ):
         super(Client, self).__init__(
             project=project, credentials=credentials, _http=_http
         )
         self.namespace = namespace
+        self._client_info = client_info
         self._batch_stack = _LocalStack()
         self._datastore_api_internal = None
         if _use_grpc is None:

datastore/tests/unit/test__gapic.py

@@ -41,7 +41,8 @@ def test_live_api(self, make_chan, mock_klass):
         client = mock.Mock(
             _base_url=base_url,
             _credentials=mock.sentinel.credentials,
-            spec=["_base_url", "_credentials"],
+            _client_info=mock.sentinel.client_info,
+            spec=["_base_url", "_credentials", "_client_info"],
         )
         ds_api = self._call_fut(client)
         self.assertIs(ds_api, mock.sentinel.ds_client)
@@ -52,7 +53,7 @@ def test_live_api(self, make_chan, mock_klass):
             "datastore.googleapis.com:443",
         )
         mock_klass.assert_called_once_with(
-            channel=mock.sentinel.channel, client_info=mock.ANY
+            channel=mock.sentinel.channel, client_info=mock.sentinel.client_info
         )
 
     @mock.patch(
@@ -70,12 +71,13 @@ def test_emulator(self, make_chan, mock_klass):
         client = mock.Mock(
             _base_url=base_url,
             _credentials=mock.sentinel.credentials,
-            spec=["_base_url", "_credentials"],
+            _client_info=mock.sentinel.client_info,
+            spec=["_base_url", "_credentials", "_client_info"],
         )
         ds_api = self._call_fut(client)
         self.assertIs(ds_api, mock.sentinel.ds_client)
 
         make_chan.assert_called_once_with(host)
         mock_klass.assert_called_once_with(
-            channel=mock.sentinel.channel, client_info=mock.ANY
+            channel=mock.sentinel.channel, client_info=mock.sentinel.client_info
         )