diff --git a/bigquery/google/cloud/bigquery/_http.py b/bigquery/google/cloud/bigquery/_http.py index 1dd7524542a5..0e5475f5f54b 100644 --- a/bigquery/google/cloud/bigquery/_http.py +++ b/bigquery/google/cloud/bigquery/_http.py @@ -14,14 +14,12 @@ """Create / interact with Google BigQuery connections.""" +import google.api_core.gapic_v1.client_info from google.cloud import _http from google.cloud.bigquery import __version__ -_CLIENT_INFO = _http.CLIENT_INFO_TEMPLATE.format(__version__) - - class Connection(_http.JSONConnection): """A connection to Google BigQuery via the JSON REST API. @@ -29,6 +27,19 @@ class Connection(_http.JSONConnection): :param client: The client that owns the current connection. """ + def __init__(self, client, client_info=None): + super(Connection, self).__init__(client) + + if client_info is None: + client_info = google.api_core.gapic_v1.client_info.ClientInfo( + gapic_version=__version__, client_library_version=__version__ + ) + else: + client_info.gapic_version = __version__ + client_info.client_library_version = __version__ + self._client_info = client_info + self._extra_headers = {} + API_BASE_URL = "https://www.googleapis.com" """The base of the API call URL.""" @@ -38,4 +49,21 @@ class Connection(_http.JSONConnection): API_URL_TEMPLATE = "{api_base_url}/bigquery/{api_version}{path}" """A template for the URL of a particular API call.""" - _EXTRA_HEADERS = {_http.CLIENT_INFO_HEADER: _CLIENT_INFO} + @property + def USER_AGENT(self): + return self._client_info.to_user_agent() + + @USER_AGENT.setter + def USER_AGENT(self, value): + self._client_info.user_agent = value + + @property + def _EXTRA_HEADERS(self): + self._extra_headers[ + _http.CLIENT_INFO_HEADER + ] = self._client_info.to_user_agent() + return self._extra_headers + + @_EXTRA_HEADERS.setter + def _EXTRA_HEADERS(self, value): + self._extra_headers = value diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index bb6a375975f2..db53dab9ef11 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -128,6 +128,11 @@ class Client(ClientWithProject): default_query_job_config (google.cloud.bigquery.job.QueryJobConfig): (Optional) Default ``QueryJobConfig``. Will be merged into job configs passed into the ``query`` method. + 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 library + or partner tool. Raises: google.auth.exceptions.DefaultCredentialsError: @@ -148,11 +153,12 @@ def __init__( _http=None, location=None, default_query_job_config=None, + client_info=None, ): super(Client, self).__init__( project=project, credentials=credentials, _http=_http ) - self._connection = Connection(self) + self._connection = Connection(self, client_info=client_info) self._location = location self._default_query_job_config = default_query_job_config diff --git a/bigquery/tests/unit/test__http.py b/bigquery/tests/unit/test__http.py index 890046ee05fa..d7d25ea445a0 100644 --- a/bigquery/tests/unit/test__http.py +++ b/bigquery/tests/unit/test__http.py @@ -45,9 +45,36 @@ def test_build_api_url_w_extra_query_params(self): parms = dict(parse_qsl(qs)) self.assertEqual(parms["bar"], "baz") + def test_user_agent(self): + from google.cloud import _http as base_http + + http = mock.create_autospec(requests.Session, instance=True) + response = requests.Response() + response.status_code = 200 + data = b"brent-spiner" + response._content = data + http.request.return_value = response + client = mock.Mock(_http=http, spec=["_http"]) + + conn = self._make_one(client) + conn.USER_AGENT = "my-application/1.2.3" + req_data = "req-data-boring" + result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False) + self.assertEqual(result, data) + + expected_headers = { + "Accept-Encoding": "gzip", + base_http.CLIENT_INFO_HEADER: conn.USER_AGENT, + "User-Agent": conn.USER_AGENT, + } + expected_uri = conn.build_api_url("/rainbow") + http.request.assert_called_once_with( + data=req_data, headers=expected_headers, method="GET", url=expected_uri + ) + self.assertIn("my-application/1.2.3", conn.USER_AGENT) + def test_extra_headers(self): from google.cloud import _http as base_http - from google.cloud.bigquery import _http as MUT http = mock.create_autospec(requests.Session, instance=True) response = requests.Response() @@ -58,14 +85,44 @@ def test_extra_headers(self): client = mock.Mock(_http=http, spec=["_http"]) conn = self._make_one(client) + conn._EXTRA_HEADERS["x-test-header"] = "a test value" + req_data = "req-data-boring" + result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False) + self.assertEqual(result, data) + + expected_headers = { + "Accept-Encoding": "gzip", + base_http.CLIENT_INFO_HEADER: conn.USER_AGENT, + "User-Agent": conn.USER_AGENT, + "x-test-header": "a test value", + } + expected_uri = conn.build_api_url("/rainbow") + http.request.assert_called_once_with( + data=req_data, headers=expected_headers, method="GET", url=expected_uri + ) + + def test_extra_headers_replace(self): + from google.cloud import _http as base_http + + http = mock.create_autospec(requests.Session, instance=True) + response = requests.Response() + response.status_code = 200 + data = b"brent-spiner" + response._content = data + http.request.return_value = response + client = mock.Mock(_http=http, spec=["_http"]) + + conn = self._make_one(client) + conn._EXTRA_HEADERS = {"x-test-header": "a test value"} req_data = "req-data-boring" result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False) self.assertEqual(result, data) expected_headers = { "Accept-Encoding": "gzip", - base_http.CLIENT_INFO_HEADER: MUT._CLIENT_INFO, + base_http.CLIENT_INFO_HEADER: conn.USER_AGENT, "User-Agent": conn.USER_AGENT, + "x-test-header": "a test value", } expected_uri = conn.build_api_url("/rainbow") http.request.assert_called_once_with( diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py index ac2e036cdb9d..08c36e0ac277 100644 --- a/bigquery/tests/unit/test_client.py +++ b/bigquery/tests/unit/test_client.py @@ -22,6 +22,7 @@ import unittest import mock +import requests import six from six.moves import http_client import pytest @@ -37,6 +38,7 @@ pyarrow = None import google.api_core.exceptions +from google.api_core.gapic_v1 import client_info import google.cloud._helpers from google.cloud.bigquery.dataset import DatasetReference @@ -1320,6 +1322,38 @@ def test_get_table(self): conn.api_request.assert_called_once_with(method="GET", path="/%s" % path) self.assertEqual(table.table_id, self.TABLE_ID) + def test_get_table_sets_user_agent(self): + creds = _make_credentials() + http = mock.create_autospec(requests.Session) + mock_response = http.request( + url=mock.ANY, method=mock.ANY, headers=mock.ANY, data=mock.ANY + ) + http.reset_mock() + mock_response.status_code = 200 + mock_response.json.return_value = self._make_table_resource() + user_agent_override = client_info.ClientInfo(user_agent="my-application/1.2.3") + client = self._make_one( + project=self.PROJECT, + credentials=creds, + client_info=user_agent_override, + _http=http, + ) + + client.get_table(self.TABLE_REF) + + expected_user_agent = user_agent_override.to_user_agent() + http.request.assert_called_once_with( + url=mock.ANY, + method="GET", + headers={ + "X-Goog-API-Client": expected_user_agent, + "Accept-Encoding": "gzip", + "User-Agent": expected_user_agent, + }, + data=mock.ANY, + ) + self.assertIn("my-application/1.2.3", expected_user_agent) + def test_update_dataset_w_invalid_field(self): from google.cloud.bigquery.dataset import Dataset