BigQuery: Set BQ storage client user-agent when in Jupyter cell

bigquery/google/cloud/bigquery/magics.py

@@ -139,15 +139,15 @@
     from IPython.core import magic_arguments
 except ImportError:  # pragma: NO COVER
     raise ImportError("This module can only be loaded in IPython.")
-try:
-    from google.cloud import bigquery_storage_v1beta1
-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
+import six
+
+
+IPYTHON_USER_AGENT = "ipython-{}".format(IPython.__version__)
 
 
 class Context(object):
@@ -399,9 +399,7 @@ 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__)
-        ),
+        client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
     )
     if context._connection:
         client._connection = context._connection
@@ -433,10 +431,24 @@ def _make_bqstorage_client(use_bqstorage_api, credentials):
     if not use_bqstorage_api:
         return None
 
-    if bigquery_storage_v1beta1 is None:
-        raise ImportError(
+    try:
+        from google.cloud import bigquery_storage_v1beta1
+    except ImportError as err:
+        customized_error = ImportError(
             "Install the google-cloud-bigquery-storage and fastavro packages "
             "to use the BigQuery Storage API."
         )
+        six.raise_from(customized_error, err)
 
-    return bigquery_storage_v1beta1.BigQueryStorageClient(credentials=credentials)
+    try:
+        from google.api_core.gapic_v1 import client_info as gapic_client_info
+    except ImportError as err:
+        customized_error = ImportError(
+            "Install the grpcio package to use the BigQuery Storage API."
+        )
+        six.raise_from(customized_error, err)
+
+    return bigquery_storage_v1beta1.BigQueryStorageClient(
+        credentials=credentials,
+        client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
+    )

bigquery/tests/unit/helpers.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import mock
+import six
+
 
 def make_connection(*responses):
     import google.cloud.bigquery._http
@@ -22,3 +25,25 @@ def make_connection(*responses):
     mock_conn.user_agent = "testing 1.2.3"
     mock_conn.api_request.side_effect = list(responses) + [NotFound("miss")]
     return mock_conn
+
+
+def maybe_fail_import(predicate):
+    """Create and return a patcher that conditionally makes an import fail.
+
+    Args:
+        predicate (Callable[[...], bool]): A callable that, if it returns `True`,
+            triggers an `ImportError`. It must accept the same arguments as the
+            built-in `__import__` function.
+            https://docs.python.org/3/library/functions.html#__import__
+
+    Returns:
+        A mock patcher object that can be used to enable patched import behavior.
+    """
+    orig_import = six.moves.builtins.__import__
+
+    def custom_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if predicate(name, globals, locals, fromlist, level):
+            raise ImportError
+        return orig_import(name, globals, locals, fromlist, level)
+
+    return mock.patch.object(six.moves.builtins, "__import__", new=custom_import)

bigquery/tests/unit/test_magics.py

@@ -42,6 +42,7 @@
 from google.cloud.bigquery import table
 from google.cloud.bigquery import magics
 from tests.unit.helpers import make_connection
+from tests.unit.helpers import maybe_fail_import
 
 
 pytestmark = pytest.mark.skipif(IPython is None, reason="Requires `ipython`")
@@ -65,6 +66,30 @@ def ipython_interactive(request, ipython):
         yield ipython
 
 
+@pytest.fixture(scope="session")
+def missing_bq_storage():
+    """Provide a patcher that can make the bigquery storage import to fail."""
+
+    def fail_if(name, globals, locals, fromlist, level):
+        # NOTE: *very* simplified, assuming a straightforward absolute import
+        return "bigquery_storage_v1beta1" in name or (
+            fromlist is not None and "bigquery_storage_v1beta1" in fromlist
+        )
+
+    return maybe_fail_import(predicate=fail_if)
+
+
+@pytest.fixture(scope="session")
+def missing_grpcio_lib():
+    """Provide a patcher that can make the gapic library import to fail."""
+
+    def fail_if(name, globals, locals, fromlist, level):
+        # NOTE: *very* simplified, assuming a straightforward absolute import
+        return "gapic_v1" in name or (fromlist is not None and "gapic_v1" in fromlist)
+
+    return maybe_fail_import(predicate=fail_if)
+
+
 JOB_REFERENCE_RESOURCE = {"projectId": "its-a-project-eh", "jobId": "some-random-id"}
 TABLE_REFERENCE_RESOURCE = {
     "projectId": "its-a-project-eh",
@@ -267,18 +292,28 @@ def test__make_bqstorage_client_true():
     assert isinstance(got, bigquery_storage_v1beta1.BigQueryStorageClient)
 
 
-def test__make_bqstorage_client_true_raises_import_error(monkeypatch):
-    monkeypatch.setattr(magics, "bigquery_storage_v1beta1", None)
+def test__make_bqstorage_client_true_raises_import_error(missing_bq_storage):
     credentials_mock = mock.create_autospec(
         google.auth.credentials.Credentials, instance=True
     )
 
-    with pytest.raises(ImportError) as exc_context:
+    with pytest.raises(ImportError) as exc_context, missing_bq_storage:
         magics._make_bqstorage_client(True, credentials_mock)
 
     assert "google-cloud-bigquery-storage" in str(exc_context.value)
 
 
+def test__make_bqstorage_client_true_missing_gapic(missing_grpcio_lib):
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    with pytest.raises(ImportError) as exc_context, missing_grpcio_lib:
+        magics._make_bqstorage_client(True, credentials_mock)
+
+    assert "grpcio" in str(exc_context.value)
+
+
 @pytest.mark.usefixtures("ipython_interactive")
 def test_extension_load():
     ip = IPython.get_ipython()
@@ -291,16 +326,13 @@ def test_extension_load():
 
 @pytest.mark.usefixtures("ipython_interactive")
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
-def test_bigquery_magic_without_optional_arguments(monkeypatch):
+def test_bigquery_magic_without_optional_arguments(missing_bq_storage):
     ip = IPython.get_ipython()
     ip.extension_manager.load_extension("google.cloud.bigquery")
     magics.context.credentials = mock.create_autospec(
         google.auth.credentials.Credentials, instance=True
     )
 
-    # Shouldn't fail when BigQuery Storage client isn't installed.
-    monkeypatch.setattr(magics, "bigquery_storage_v1beta1", None)
-
     sql = "SELECT 17 AS num"
     result = pandas.DataFrame([17], columns=["num"])
     run_query_patch = mock.patch(
@@ -310,9 +342,10 @@ def test_bigquery_magic_without_optional_arguments(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
-    with run_query_patch as run_query_mock:
-        run_query_mock.return_value = query_job_mock
 
+    # Shouldn't fail when BigQuery Storage client isn't installed.
+    with run_query_patch as run_query_mock, missing_bq_storage:
+        run_query_mock.return_value = query_job_mock
         return_value = ip.run_cell_magic("bigquery", "", sql)
 
     assert isinstance(return_value, pandas.DataFrame)
@@ -459,8 +492,8 @@ def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch):
         bigquery_storage_v1beta1.BigQueryStorageClient, instance=True
     )
     bqstorage_mock.return_value = bqstorage_instance_mock
-    monkeypatch.setattr(
-        magics.bigquery_storage_v1beta1, "BigQueryStorageClient", bqstorage_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient", bqstorage_mock
     )
 
     sql = "SELECT 17 AS num"
@@ -472,15 +505,21 @@ def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
-    with run_query_patch as run_query_mock:
+    with run_query_patch as run_query_mock, bqstorage_client_patch:
         run_query_mock.return_value = query_job_mock
 
         return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api", sql)
 
-        bqstorage_mock.assert_called_once_with(credentials=mock_credentials)
-        query_job_mock.to_dataframe.assert_called_once_with(
-            bqstorage_client=bqstorage_instance_mock
-        )
+    assert len(bqstorage_mock.call_args_list) == 1
+    kwargs = bqstorage_mock.call_args_list[0].kwargs
+    assert kwargs.get("credentials") is mock_credentials
+    client_info = kwargs.get("client_info")
+    assert client_info is not None
+    assert client_info.user_agent == "ipython-" + IPython.__version__
+
+    query_job_mock.to_dataframe.assert_called_once_with(
+        bqstorage_client=bqstorage_instance_mock
+    )
 
     assert isinstance(return_value, pandas.DataFrame)
 
@@ -509,8 +548,8 @@ def test_bigquery_magic_with_bqstorage_from_context(monkeypatch):
         bigquery_storage_v1beta1.BigQueryStorageClient, instance=True
     )
     bqstorage_mock.return_value = bqstorage_instance_mock
-    monkeypatch.setattr(
-        magics.bigquery_storage_v1beta1, "BigQueryStorageClient", bqstorage_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient", bqstorage_mock
     )
 
     sql = "SELECT 17 AS num"
@@ -522,15 +561,21 @@ def test_bigquery_magic_with_bqstorage_from_context(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
-    with run_query_patch as run_query_mock:
+    with run_query_patch as run_query_mock, bqstorage_client_patch:
         run_query_mock.return_value = query_job_mock
 
         return_value = ip.run_cell_magic("bigquery", "", sql)
 
-        bqstorage_mock.assert_called_once_with(credentials=mock_credentials)
-        query_job_mock.to_dataframe.assert_called_once_with(
-            bqstorage_client=bqstorage_instance_mock
-        )
+    assert len(bqstorage_mock.call_args_list) == 1
+    kwargs = bqstorage_mock.call_args_list[0].kwargs
+    assert kwargs.get("credentials") is mock_credentials
+    client_info = kwargs.get("client_info")
+    assert client_info is not None
+    assert client_info.user_agent == "ipython-" + IPython.__version__
+
+    query_job_mock.to_dataframe.assert_called_once_with(
+        bqstorage_client=bqstorage_instance_mock
+    )
 
     assert isinstance(return_value, pandas.DataFrame)
 
@@ -554,8 +599,8 @@ def test_bigquery_magic_without_bqstorage(monkeypatch):
     bqstorage_mock = mock.create_autospec(
         bigquery_storage_v1beta1.BigQueryStorageClient
     )
-    monkeypatch.setattr(
-        magics.bigquery_storage_v1beta1, "BigQueryStorageClient", bqstorage_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient", bqstorage_mock
     )
 
     sql = "SELECT 17 AS num"
@@ -567,7 +612,7 @@ def test_bigquery_magic_without_bqstorage(monkeypatch):
         google.cloud.bigquery.job.QueryJob, instance=True
     )
     query_job_mock.to_dataframe.return_value = result
-    with run_query_patch as run_query_mock:
+    with run_query_patch as run_query_mock, bqstorage_client_patch:
         run_query_mock.return_value = query_job_mock
 
         return_value = ip.run_cell_magic("bigquery", "", sql)