Skip to content

Commit b91d86a

Browse files
plamuttswast
authored andcommitted
Set IPython user agent when running queries with IPython cell magic (#8713)
* Set ipython user agent when in IPython cell * Revert year bump in license headers * Fix false positive failure in pyarrow warning test Any additional warnings on top of the expected pyarrow warning would cause one of the tests to fail. This commit makes that test more robust by only focusing on the existence of the specific warning of interest. * Modify test - *all* warnings should be pyarrow's
1 parent f5ad704 commit b91d86a

3 files changed

Lines changed: 33 additions & 4 deletions

File tree

‎bigquery/google/cloud/bigquery/magics.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@
144144
except ImportError: # pragma: NO COVER
145145
bigquery_storage_v1beta1 = None
146146

147+
from google.api_core import client_info
147148
import google.auth
148149
from google.cloud import bigquery
149150
from google.cloud.bigquery.dbapi import _helpers
@@ -398,6 +399,9 @@ def _cell_magic(line, query):
398399
project=project,
399400
credentials=context.credentials,
400401
default_query_job_config=context.default_query_job_config,
402+
client_info=client_info.ClientInfo(
403+
user_agent="ipython-{}".format(IPython.__version__)
404+
),
401405
)
402406
if context._connection:
403407
client._connection = context._connection

‎bigquery/tests/unit/test_client.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5353,10 +5353,10 @@ def test_load_table_from_dataframe_w_schema_wo_pyarrow(self):
53535353
dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION
53545354
)
53555355

5356-
assert len(warned) == 1
5357-
warning = warned[0]
5358-
assert warning.category is PendingDeprecationWarning
5359-
assert "pyarrow" in str(warning)
5356+
assert warned # there should be at least one warning
5357+
for warning in warned:
5358+
assert "pyarrow" in str(warning)
5359+
assert warning.category in (DeprecationWarning, PendingDeprecationWarning)
53605360

53615361
load_table_from_file.assert_called_once_with(
53625362
client,

‎bigquery/tests/unit/test_magics.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,31 @@ def test_bigquery_magic_without_optional_arguments(monkeypatch):
320320
assert list(return_value) == list(result) # verify column names
321321

322322

323+
@pytest.mark.usefixtures("ipython_interactive")
324+
def test_bigquery_magic_default_connection_user_agent():
325+
ip = IPython.get_ipython()
326+
ip.extension_manager.load_extension("google.cloud.bigquery")
327+
magics.context._connection = None
328+
329+
credentials_mock = mock.create_autospec(
330+
google.auth.credentials.Credentials, instance=True
331+
)
332+
default_patch = mock.patch(
333+
"google.auth.default", return_value=(credentials_mock, "general-project")
334+
)
335+
run_query_patch = mock.patch(
336+
"google.cloud.bigquery.magics._run_query", autospec=True
337+
)
338+
conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
339+
340+
with conn_patch as conn, run_query_patch, default_patch:
341+
ip.run_cell_magic("bigquery", "", "SELECT 17 as num")
342+
343+
client_info_arg = conn.call_args.kwargs.get("client_info")
344+
assert client_info_arg is not None
345+
assert client_info_arg.user_agent == "ipython-" + IPython.__version__
346+
347+
323348
@pytest.mark.usefixtures("ipython_interactive")
324349
def test_bigquery_magic_with_legacy_sql():
325350
ip = IPython.get_ipython()

0 commit comments

Comments
 (0)