Skip to content

Commit 0a55c26

Browse files
tswasttswast
authored
Add client_info to BigQuery constructor for user-amenable user agent headers (#7806)
* Add client_info to BigQuery constructor for user-amenable user agent headers This aligns BigQuery's behavior regarding the User-Agent and X-Goog-Api-Client headers with that of the GAPIC-based clients. Old: X-Goog-API-Client: gl-python/3.7.2 gccl/1.11.2 User-Agent: gcloud-python/0.29.1 New: X-Goog-API-Client: optional-application-id/1.2.3 gl-python/3.7.2 grpc/1.20.0 gax/1.9.0 gapic/1.11.2 gccl/1.11.2 User-Agent: optional-application-id/1.2.3 gl-python/3.7.2 grpc/1.20.0 gax/1.9.0 gapic/1.11.2 gccl/1.11.2 In order to set the `optional-application-id/1.2.3`, the latest version of `api_core` is required, but since that's an uncommon usecase and it doesn't break, just ignore the custom User-Agent if an older version is used, I didn't update the minimum version `setup.py`. * flake8 & blacken
1 parent 85e8012 commit 0a55c26

File tree

4 files changed

+132
-7
lines changed

4 files changed

+132
-7
lines changed

‎bigquery/google/cloud/bigquery/_http.py‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,21 +14,32 @@
1414

1515
"""Create / interact with Google BigQuery connections."""
1616

17+
import google.api_core.gapic_v1.client_info
1718
from google.cloud import _http
1819

1920
from google.cloud.bigquery import __version__
2021

2122

22-
_CLIENT_INFO = _http.CLIENT_INFO_TEMPLATE.format(__version__)
23-
24-
2523
class Connection(_http.JSONConnection):
2624
"""A connection to Google BigQuery via the JSON REST API.
2725
2826
:type client: :class:`~google.cloud.bigquery.client.Client`
2927
:param client: The client that owns the current connection.
3028
"""
3129

30+
def __init__(self, client, client_info=None):
31+
super(Connection, self).__init__(client)
32+
33+
if client_info is None:
34+
client_info = google.api_core.gapic_v1.client_info.ClientInfo(
35+
gapic_version=__version__, client_library_version=__version__
36+
)
37+
else:
38+
client_info.gapic_version = __version__
39+
client_info.client_library_version = __version__
40+
self._client_info = client_info
41+
self._extra_headers = {}
42+
3243
API_BASE_URL = "https://www.googleapis.com"
3344
"""The base of the API call URL."""
3445

@@ -38,4 +49,21 @@ class Connection(_http.JSONConnection):
3849
API_URL_TEMPLATE = "{api_base_url}/bigquery/{api_version}{path}"
3950
"""A template for the URL of a particular API call."""
4051

41-
_EXTRA_HEADERS = {_http.CLIENT_INFO_HEADER: _CLIENT_INFO}
52+
@property
53+
def USER_AGENT(self):
54+
return self._client_info.to_user_agent()
55+
56+
@USER_AGENT.setter
57+
def USER_AGENT(self, value):
58+
self._client_info.user_agent = value
59+
60+
@property
61+
def _EXTRA_HEADERS(self):
62+
self._extra_headers[
63+
_http.CLIENT_INFO_HEADER
64+
] = self._client_info.to_user_agent()
65+
return self._extra_headers
66+
67+
@_EXTRA_HEADERS.setter
68+
def _EXTRA_HEADERS(self, value):
69+
self._extra_headers = value

‎bigquery/google/cloud/bigquery/client.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,11 @@ class Client(ClientWithProject):
128128
default_query_job_config (google.cloud.bigquery.job.QueryJobConfig):
129129
(Optional) Default ``QueryJobConfig``.
130130
Will be merged into job configs passed into the ``query`` method.
131+
client_info (google.api_core.gapic_v1.client_info.ClientInfo):
132+
The client info used to send a user-agent string along with API
133+
requests. If ``None``, then default info will be used. Generally,
134+
you only need to set this if you're developing your own library
135+
or partner tool.
131136
132137
Raises:
133138
google.auth.exceptions.DefaultCredentialsError:
@@ -148,11 +153,12 @@ def __init__(
148153
_http=None,
149154
location=None,
150155
default_query_job_config=None,
156+
client_info=None,
151157
):
152158
super(Client, self).__init__(
153159
project=project, credentials=credentials, _http=_http
154160
)
155-
self._connection = Connection(self)
161+
self._connection = Connection(self, client_info=client_info)
156162
self._location = location
157163
self._default_query_job_config = default_query_job_config
158164

‎bigquery/tests/unit/test__http.py‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,36 @@ def test_build_api_url_w_extra_query_params(self):
4545
parms = dict(parse_qsl(qs))
4646
self.assertEqual(parms["bar"], "baz")
4747

48+
def test_user_agent(self):
49+
from google.cloud import _http as base_http
50+
51+
http = mock.create_autospec(requests.Session, instance=True)
52+
response = requests.Response()
53+
response.status_code = 200
54+
data = b"brent-spiner"
55+
response._content = data
56+
http.request.return_value = response
57+
client = mock.Mock(_http=http, spec=["_http"])
58+
59+
conn = self._make_one(client)
60+
conn.USER_AGENT = "my-application/1.2.3"
61+
req_data = "req-data-boring"
62+
result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False)
63+
self.assertEqual(result, data)
64+
65+
expected_headers = {
66+
"Accept-Encoding": "gzip",
67+
base_http.CLIENT_INFO_HEADER: conn.USER_AGENT,
68+
"User-Agent": conn.USER_AGENT,
69+
}
70+
expected_uri = conn.build_api_url("/rainbow")
71+
http.request.assert_called_once_with(
72+
data=req_data, headers=expected_headers, method="GET", url=expected_uri
73+
)
74+
self.assertIn("my-application/1.2.3", conn.USER_AGENT)
75+
4876
def test_extra_headers(self):
4977
from google.cloud import _http as base_http
50-
from google.cloud.bigquery import _http as MUT
5178

5279
http = mock.create_autospec(requests.Session, instance=True)
5380
response = requests.Response()
@@ -58,14 +85,44 @@ def test_extra_headers(self):
5885
client = mock.Mock(_http=http, spec=["_http"])
5986

6087
conn = self._make_one(client)
88+
conn._EXTRA_HEADERS["x-test-header"] = "a test value"
89+
req_data = "req-data-boring"
90+
result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False)
91+
self.assertEqual(result, data)
92+
93+
expected_headers = {
94+
"Accept-Encoding": "gzip",
95+
base_http.CLIENT_INFO_HEADER: conn.USER_AGENT,
96+
"User-Agent": conn.USER_AGENT,
97+
"x-test-header": "a test value",
98+
}
99+
expected_uri = conn.build_api_url("/rainbow")
100+
http.request.assert_called_once_with(
101+
data=req_data, headers=expected_headers, method="GET", url=expected_uri
102+
)
103+
104+
def test_extra_headers_replace(self):
105+
from google.cloud import _http as base_http
106+
107+
http = mock.create_autospec(requests.Session, instance=True)
108+
response = requests.Response()
109+
response.status_code = 200
110+
data = b"brent-spiner"
111+
response._content = data
112+
http.request.return_value = response
113+
client = mock.Mock(_http=http, spec=["_http"])
114+
115+
conn = self._make_one(client)
116+
conn._EXTRA_HEADERS = {"x-test-header": "a test value"}
61117
req_data = "req-data-boring"
62118
result = conn.api_request("GET", "/rainbow", data=req_data, expect_json=False)
63119
self.assertEqual(result, data)
64120

65121
expected_headers = {
66122
"Accept-Encoding": "gzip",
67-
base_http.CLIENT_INFO_HEADER: MUT._CLIENT_INFO,
123+
base_http.CLIENT_INFO_HEADER: conn.USER_AGENT,
68124
"User-Agent": conn.USER_AGENT,
125+
"x-test-header": "a test value",
69126
}
70127
expected_uri = conn.build_api_url("/rainbow")
71128
http.request.assert_called_once_with(

‎bigquery/tests/unit/test_client.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import unittest
2323

2424
import mock
25+
import requests
2526
import six
2627
from six.moves import http_client
2728
import pytest
@@ -37,6 +38,7 @@
3738
pyarrow = None
3839

3940
import google.api_core.exceptions
41+
from google.api_core.gapic_v1 import client_info
4042
import google.cloud._helpers
4143
from google.cloud.bigquery.dataset import DatasetReference
4244

@@ -1320,6 +1322,38 @@ def test_get_table(self):
13201322
conn.api_request.assert_called_once_with(method="GET", path="/%s" % path)
13211323
self.assertEqual(table.table_id, self.TABLE_ID)
13221324

1325+
def test_get_table_sets_user_agent(self):
1326+
creds = _make_credentials()
1327+
http = mock.create_autospec(requests.Session)
1328+
mock_response = http.request(
1329+
url=mock.ANY, method=mock.ANY, headers=mock.ANY, data=mock.ANY
1330+
)
1331+
http.reset_mock()
1332+
mock_response.status_code = 200
1333+
mock_response.json.return_value = self._make_table_resource()
1334+
user_agent_override = client_info.ClientInfo(user_agent="my-application/1.2.3")
1335+
client = self._make_one(
1336+
project=self.PROJECT,
1337+
credentials=creds,
1338+
client_info=user_agent_override,
1339+
_http=http,
1340+
)
1341+
1342+
client.get_table(self.TABLE_REF)
1343+
1344+
expected_user_agent = user_agent_override.to_user_agent()
1345+
http.request.assert_called_once_with(
1346+
url=mock.ANY,
1347+
method="GET",
1348+
headers={
1349+
"X-Goog-API-Client": expected_user_agent,
1350+
"Accept-Encoding": "gzip",
1351+
"User-Agent": expected_user_agent,
1352+
},
1353+
data=mock.ANY,
1354+
)
1355+
self.assertIn("my-application/1.2.3", expected_user_agent)
1356+
13231357
def test_update_dataset_w_invalid_field(self):
13241358
from google.cloud.bigquery.dataset import Dataset
13251359

0 commit comments

Comments
 (0)