From 3ea0456229ac68df175ac10ee2da3a48c409143c Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Tue, 23 Jul 2019 10:33:17 +0200 Subject: [PATCH 1/5] Move optional bqstorage import into function --- bigquery/google/cloud/bigquery/magics.py | 12 +++--- bigquery/tests/unit/test_magics.py | 49 +++++++++++++++--------- 2 files changed, 37 insertions(+), 24 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 9bf2019c5c2e..76eae1d06c6b 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -139,15 +139,12 @@ 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 class Context(object): @@ -433,10 +430,13 @@ 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) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 44e0571d1ee4..bff1e6d1c8c1 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -65,6 +65,22 @@ 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.""" + orig_import = six.moves.builtins.__import__ + + def custom_import(name, globals=None, locals=None, fromlist=(), level=0): + # NOTE: *very* simplified, assuming a straightforward absolute import + if "bigquery_storage_v1beta1" in name or ( + fromlist is not None and "bigquery_storage_v1beta1" in fromlist + ): + raise ImportError + return orig_import(name, globals, locals, fromlist, level) + + return mock.patch.object(six.moves.builtins, "__import__", new=custom_import) + + JOB_REFERENCE_RESOURCE = {"projectId": "its-a-project-eh", "jobId": "some-random-id"} TABLE_REFERENCE_RESOURCE = { "projectId": "its-a-project-eh", @@ -267,13 +283,12 @@ 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) @@ -291,16 +306,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 +322,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 +472,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,7 +485,7 @@ 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) @@ -509,8 +522,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,7 +535,7 @@ 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) @@ -554,8 +567,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 +580,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) From b323b206cd5b0d1be3937162d9583a2b3ae4d14c Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Tue, 23 Jul 2019 15:03:35 +0200 Subject: [PATCH 2/5] Set BQ storage client useragent in IPython cell --- bigquery/google/cloud/bigquery/magics.py | 12 ++++++---- bigquery/tests/unit/test_magics.py | 28 +++++++++++++++++------- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 76eae1d06c6b..bd31812e7131 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -147,6 +147,9 @@ import six +IPYTHON_USER_AGENT = "ipython-{}".format(IPython.__version__) + + class Context(object): """Storage for objects to be used throughout an IPython notebook session. @@ -396,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 @@ -439,4 +440,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials): ) six.raise_from(customized_error, err) - return bigquery_storage_v1beta1.BigQueryStorageClient(credentials=credentials) + return bigquery_storage_v1beta1.BigQueryStorageClient( + credentials=credentials, + client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), + ) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index bff1e6d1c8c1..1bfbcee8bc56 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -490,10 +490,16 @@ def test_bigquery_magic_with_bqstorage_from_argument(monkeypatch): 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) @@ -540,10 +546,16 @@ def test_bigquery_magic_with_bqstorage_from_context(monkeypatch): 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) From 1329f2da582a8e243469c50783c640bff2e1f53c Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Fri, 26 Jul 2019 14:49:03 +0200 Subject: [PATCH 3/5] Extract import patcher factory to test helpers --- bigquery/tests/unit/helpers.py | 25 +++++++++++++++++++++++++ bigquery/tests/unit/test_magics.py | 12 +++++------- 2 files changed, 30 insertions(+), 7 deletions(-) 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 1bfbcee8bc56..f591bd623309 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`") @@ -68,17 +69,14 @@ def ipython_interactive(request, ipython): @pytest.fixture(scope="session") def missing_bq_storage(): """Provide a patcher that can make the bigquery storage import to fail.""" - orig_import = six.moves.builtins.__import__ - def custom_import(name, globals=None, locals=None, fromlist=(), level=0): + def fail_if(name, globals, locals, fromlist, level): # NOTE: *very* simplified, assuming a straightforward absolute import - if "bigquery_storage_v1beta1" in name or ( + return "bigquery_storage_v1beta1" in name or ( fromlist is not None and "bigquery_storage_v1beta1" in fromlist - ): - raise ImportError - return orig_import(name, globals, locals, fromlist, level) + ) - return mock.patch.object(six.moves.builtins, "__import__", new=custom_import) + return maybe_fail_import(predicate=fail_if) JOB_REFERENCE_RESOURCE = {"projectId": "its-a-project-eh", "jobId": "some-random-id"} From 0d71c5c7d45a22f6fe434a4406f83d5266b261f2 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Mon, 29 Jul 2019 12:42:50 +0200 Subject: [PATCH 4/5] Use ClientInfo from gapic for BQ storage client The client is a grpc client, thus the grpc ClientInfo class should be used to configure it. --- bigquery/google/cloud/bigquery/magics.py | 10 +++++++++- bigquery/tests/unit/test_magics.py | 22 ++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index bd31812e7131..d767db83a4e1 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -440,7 +440,15 @@ def _make_bqstorage_client(use_bqstorage_api, credentials): ) six.raise_from(customized_error, err) + 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=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), + client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), ) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index f591bd623309..52d971d3cc59 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -79,6 +79,17 @@ def fail_if(name, globals, locals, fromlist, level): 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", @@ -292,6 +303,17 @@ def test__make_bqstorage_client_true_raises_import_error(missing_bq_storage): 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() From 4414a318bf1a35c85b02b4df8d6753b1fc8a31ea Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Tue, 30 Jul 2019 08:21:12 +0200 Subject: [PATCH 5/5] Replace fastavro with pyerror in import error msg --- bigquery/google/cloud/bigquery/magics.py | 2 +- bigquery/tests/unit/test_magics.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index d767db83a4e1..44596f2ef88e 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -435,7 +435,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials): from google.cloud import bigquery_storage_v1beta1 except ImportError as err: customized_error = ImportError( - "Install the google-cloud-bigquery-storage and fastavro packages " + "Install the google-cloud-bigquery-storage and pyarrow packages " "to use the BigQuery Storage API." ) six.raise_from(customized_error, err) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 52d971d3cc59..760b6ccf568d 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -300,7 +300,9 @@ def test__make_bqstorage_client_true_raises_import_error(missing_bq_storage): 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) + 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):