diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 0acde4f21b5f..9bf2019c5c2e 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -144,6 +144,7 @@ except ImportError: # pragma: NO COVER bigquery_storage_v1beta1 = None +from google.api_core import client_info import google.auth from google.cloud import bigquery from google.cloud.bigquery.dbapi import _helpers @@ -398,6 +399,9 @@ def _cell_magic(line, query): project=project, credentials=context.credentials, default_query_job_config=context.default_query_job_config, + client_info=client_info.ClientInfo( + user_agent="ipython-{}".format(IPython.__version__) + ), ) if context._connection: client._connection = context._connection diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py index 8ad9dc8858c6..2be40a52e1fc 100644 --- a/bigquery/tests/unit/test_client.py +++ b/bigquery/tests/unit/test_client.py @@ -5353,10 +5353,10 @@ def test_load_table_from_dataframe_w_schema_wo_pyarrow(self): dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION ) - assert len(warned) == 1 - warning = warned[0] - assert warning.category is PendingDeprecationWarning - assert "pyarrow" in str(warning) + assert warned # there should be at least one warning + for warning in warned: + assert "pyarrow" in str(warning) + assert warning.category in (DeprecationWarning, PendingDeprecationWarning) load_table_from_file.assert_called_once_with( client, diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 6e8c941bbc02..1b2a4dd5932f 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -320,6 +320,31 @@ def test_bigquery_magic_without_optional_arguments(monkeypatch): assert list(return_value) == list(result) # verify column names +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_default_connection_user_agent(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._connection = None + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True) + + with conn_patch as conn, run_query_patch, default_patch: + ip.run_cell_magic("bigquery", "", "SELECT 17 as num") + + client_info_arg = conn.call_args.kwargs.get("client_info") + assert client_info_arg is not None + assert client_info_arg.user_agent == "ipython-" + IPython.__version__ + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_with_legacy_sql(): ip = IPython.get_ipython()