8000 Add `client_info` to BigQuery constructor for user-amenable user agent headers by tswast · Pull Request #7806 · googleapis/google-cloud-python · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 32 additions & 4 deletions bigquery/google/cloud/bigquery/_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,32 @@

"""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.

:type client: :class:`~google.cloud.bigquery.client.Client`
:param client: The client that owns the current connection.
"""

def __init__(self, client, client_info=None):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tswast Need to document the client_info parameter in the class docstring.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should probably hoist this machinery into the google.cloud._http.JSONConnection base class.

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."""

Expand All @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These aren't constants, so why uppercase?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For backwards compatibility.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's a USER_AGENT on the superclass, which we want to override the behavior of.


@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
8 changes: 7 additions & 1 deletion bigquery/google/cloud/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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

Expand Down
61 changes: 59 additions & 2 deletions bigquery/tests/unit/test__http.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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(
Expand Down
A36C
34 changes: 34 additions & 0 deletions bigquery/tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import unittest

import mock
import requests
import six
from six.moves import http_client
import pytest
Expand All @@ -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

Expand Down Expand Up @@ -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

Expand Down
0