diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 9bf2019c5c2e..44596f2ef88e 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/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( - "Install the google-cloud-bigquery-storage and fastavro packages " + try: + from google.cloud import bigquery_storage_v1beta1 + except ImportError as err: + customized_error = ImportError( + "Install the google-cloud-bigquery-storage and pyarrow 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), + ) diff --git a/bigquery/tests/unit/helpers.py b/bigquery/tests/unit/helpers.py index 5b731a763a99..673aa8ac5f02 100644 --- a/bigquery/tests/unit/helpers.py +++ b/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) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 44e0571d1ee4..760b6ccf568d 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/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,16 +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, missing_bq_storage: + magics._make_bqstorage_client(True, credentials_mock) + + error_msg = str(exc_context.value) + assert "google-cloud-bigquery-storage" in error_msg + assert "pyarrow" in error_msg + + +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: + with pytest.raises(ImportError) as exc_context, missing_grpcio_lib: magics._make_bqstorage_client(True, credentials_mock) - assert "google-cloud-bigquery-storage" in str(exc_context.value) + assert "grpcio" in str(exc_context.value) @pytest.mark.usefixtures("ipython_interactive") @@ -291,16 +328,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 +344,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 +494,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 +507,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 +550,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 +563,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 +601,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 +614,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)