From d70a5138c4a85a23da6855f82b599823067fc048 Mon Sep 17 00:00:00 2001 From: Carlos de la Guardia Date: Tue, 13 Oct 2020 02:06:35 -0500 Subject: [PATCH 1/3] feat: allow client options to be set in magics context --- google/cloud/bigquery/magics/magics.py | 49 ++++++++++++++++++++++++-- tests/unit/test_magics.py | 41 +++++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index 22175ee45..dc224bbc8 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -139,6 +139,7 @@ import re import ast +import copy import functools import sys import time @@ -155,6 +156,7 @@ import six from google.api_core import client_info +from google.api_core import client_options from google.api_core.exceptions import NotFound import google.auth from google.cloud import bigquery @@ -178,11 +180,12 @@ def __init__(self): self._project = None self._connection = None self._default_query_job_config = bigquery.QueryJobConfig() + self._client_options = client_options.ClientOptions() @property def credentials(self): """google.auth.credentials.Credentials: Credentials to use for queries - performed through IPython magics + performed through IPython magics. Note: These credentials do not need to be explicitly defined if you are @@ -217,7 +220,7 @@ def credentials(self, value): @property def project(self): """str: Default project to use for queries performed through IPython - magics + magics. Note: The project does not need to be explicitly defined if you have an @@ -239,6 +242,30 @@ def project(self): def project(self, value): self._project = value + @property + def client_options(self): + """google.api_core.client_options.ClientOptions: client options to be + used through IPython magics. + + Note:: + The client options do not need to be explicitly defined if no + special network connections are required. Normally you would be + using the https://www.googleapis.com/ end point. + + Example: + Manually setting the endpoint: + + >>> from google.cloud.bigquery import magics + >>> client_options = {} + >>> client_options['api_endpoint'] = "https://some.special.url" + >>> magics.context.client_options = client_options + """ + return self._client_options + + @client_options.setter + def client_options(self, value): + self._client_options = value + @property def default_query_job_config(self): """google.cloud.bigquery.job.QueryJobConfig: Default job @@ -410,6 +437,15 @@ def _create_dataset_if_necessary(client, dataset_id): "Standard SQL if this argument is not used." ), ) +@magic_arguments.argument( + "--api_endpoint", + type=str, + default=None, + help=( + "The desired API endpoint, e.g., compute.googlepis.com. Defaults to this " + "option's value in the context client options." + ), +) @magic_arguments.argument( "--use_bqstorage_api", action="store_true", @@ -511,11 +547,20 @@ def _cell_magic(line, query): params = _helpers.to_query_parameters(ast.literal_eval(params_option_value)) project = args.project or context.project + + client_options = copy.deepcopy(context.client_options) + if args.api_endpoint: + if isinstance(client_options, dict): + client_options["api_endpoint"] = args.api_endpoint + else: + client_options.api_endpoint = args.api_endpoint + client = bigquery.Client( project=project, credentials=context.credentials, default_query_job_config=context.default_query_job_config, client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), + client_options=client_options, ) if context._connection: client._connection = context._connection diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py index 20be6b755..a96914f60 100644 --- a/tests/unit/test_magics.py +++ b/tests/unit/test_magics.py @@ -1180,6 +1180,47 @@ def test_bigquery_magic_with_project(): assert magics.context.project == "general-project" +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_with_api_endpoint(ipython_ns_cleanup): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._connection = None + + run_query_patch = mock.patch( + "google.cloud.bigquery.magics.magics._run_query", autospec=True + ) + with run_query_patch as run_query_mock: + ip.run_cell_magic( + "bigquery", "--api_endpoint=https://api.endpoint.com", "SELECT 17 as num" + ) + + connection_used = run_query_mock.call_args_list[0][0][0]._connection + assert connection_used.API_BASE_URL == "https://api.endpoint.com" + # context client options should not change + assert magics.context.client_options.api_endpoint is None + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_with_api_endpoint_context_dict(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._connection = None + magics.context.client_options = {} + + run_query_patch = mock.patch( + "google.cloud.bigquery.magics.magics._run_query", autospec=True + ) + with run_query_patch as run_query_mock: + ip.run_cell_magic( + "bigquery", "--api_endpoint=https://api.endpoint.com", "SELECT 17 as num" + ) + + connection_used = run_query_mock.call_args_list[0][0][0]._connection + assert connection_used.API_BASE_URL == "https://api.endpoint.com" + # context client options should not change + assert magics.context.client_options == {} + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_with_multiple_options(): ip = IPython.get_ipython() From 74802fceb70a6b7826a6257aada0c928ee917f61 Mon Sep 17 00:00:00 2001 From: Carlos de la Guardia Date: Tue, 13 Oct 2020 22:21:51 -0500 Subject: [PATCH 2/3] add separate client options for storage client --- google/cloud/bigquery/magics/magics.py | 83 ++++++++++++++++++++------ tests/unit/test_magics.py | 75 +++++++++++++++++++---- 2 files changed, 126 insertions(+), 32 deletions(-) diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index dc224bbc8..1f2ee54c7 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -180,7 +180,8 @@ def __init__(self): self._project = None self._connection = None self._default_query_job_config = bigquery.QueryJobConfig() - self._client_options = client_options.ClientOptions() + self._bigquery_client_options = client_options.ClientOptions() + self._storage_client_options = client_options.ClientOptions() @property def credentials(self): @@ -243,14 +244,14 @@ def project(self, value): self._project = value @property - def client_options(self): + def bigquery_client_options(self): """google.api_core.client_options.ClientOptions: client options to be used through IPython magics. Note:: The client options do not need to be explicitly defined if no special network connections are required. Normally you would be - using the https://www.googleapis.com/ end point. + using the https://bigquery.googleapis.com/ end point. Example: Manually setting the endpoint: @@ -258,13 +259,37 @@ def client_options(self): >>> from google.cloud.bigquery import magics >>> client_options = {} >>> client_options['api_endpoint'] = "https://some.special.url" - >>> magics.context.client_options = client_options + >>> magics.context.bigquery_client_options = client_options """ - return self._client_options + return self._bigquery_client_options - @client_options.setter - def client_options(self, value): - self._client_options = value + @bigquery_client_options.setter + def bigquery_client_options(self, value): + self._bigquery_client_options = value + + @property + def storage_client_options(self): + """google.api_core.client_options.ClientOptions: client options to be + used through IPython magics for the storage client. + + Note:: + The client options do not need to be explicitly defined if no + special network connections are required. Normally you would be + using the https://bigquerystorage.googleapis.com/ end point. + + Example: + Manually setting the endpoint: + + >>> from google.cloud.bigquery import magics + >>> client_options = {} + >>> client_options['api_endpoint'] = "https://some.special.url" + >>> magics.context.storage_client_options = client_options + """ + return self._storage_client_options + + @storage_client_options.setter + def storage_client_options(self, value): + self._storage_client_options = value @property def default_query_job_config(self): @@ -438,12 +463,21 @@ def _create_dataset_if_necessary(client, dataset_id): ), ) @magic_arguments.argument( - "--api_endpoint", + "--bigquery_api_endpoint", + type=str, + default=None, + help=( + "The desired API endpoint, e.g., bigquery.googlepis.com. Defaults to this " + "option's value in the context bigquery_client_options." + ), +) +@magic_arguments.argument( + "--storage_api_endpoint", type=str, default=None, help=( - "The desired API endpoint, e.g., compute.googlepis.com. Defaults to this " - "option's value in the context client options." + "The desired API endpoint, e.g., bigquery.googlepis.com. Defaults to this " + "option's value in the context storage_client_options." ), ) @magic_arguments.argument( @@ -548,23 +582,33 @@ def _cell_magic(line, query): project = args.project or context.project - client_options = copy.deepcopy(context.client_options) - if args.api_endpoint: - if isinstance(client_options, dict): - client_options["api_endpoint"] = args.api_endpoint + bigquery_client_options = copy.deepcopy(context.bigquery_client_options) + if args.bigquery_api_endpoint: + if isinstance(bigquery_client_options, dict): + bigquery_client_options["api_endpoint"] = args.bigquery_api_endpoint else: - client_options.api_endpoint = args.api_endpoint + bigquery_client_options.api_endpoint = args.bigquery_api_endpoint client = bigquery.Client( project=project, credentials=context.credentials, default_query_job_config=context.default_query_job_config, client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), - client_options=client_options, + client_options=bigquery_client_options, ) if context._connection: client._connection = context._connection - bqstorage_client = _make_bqstorage_client(use_bqstorage_api, context.credentials) + + storage_client_options = copy.deepcopy(context.storage_client_options) + if args.storage_api_endpoint: + if isinstance(storage_client_options, dict): + storage_client_options["api_endpoint"] = args.storage_api_endpoint + else: + storage_client_options.api_endpoint = args.storage_api_endpoint + + bqstorage_client = _make_bqstorage_client( + use_bqstorage_api, context.credentials, storage_client_options, + ) close_transports = functools.partial(_close_transports, client, bqstorage_client) @@ -677,7 +721,7 @@ def _split_args_line(line): return params_option_value, rest_of_args -def _make_bqstorage_client(use_bqstorage_api, credentials): +def _make_bqstorage_client(use_bqstorage_api, credentials, client_options): if not use_bqstorage_api: return None @@ -703,6 +747,7 @@ def _make_bqstorage_client(use_bqstorage_api, credentials): return bigquery_storage.BigQueryReadClient( credentials=credentials, client_info=gapic_client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT), + client_options=client_options, ) diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py index a96914f60..a3eec636d 100644 --- a/tests/unit/test_magics.py +++ b/tests/unit/test_magics.py @@ -309,7 +309,7 @@ def test__make_bqstorage_client_false(): credentials_mock = mock.create_autospec( google.auth.credentials.Credentials, instance=True ) - got = magics._make_bqstorage_client(False, credentials_mock) + got = magics._make_bqstorage_client(False, credentials_mock, {}) assert got is None @@ -320,7 +320,7 @@ def test__make_bqstorage_client_true(): credentials_mock = mock.create_autospec( google.auth.credentials.Credentials, instance=True ) - got = magics._make_bqstorage_client(True, credentials_mock) + got = magics._make_bqstorage_client(True, credentials_mock, {}) assert isinstance(got, bigquery_storage.BigQueryReadClient) @@ -330,7 +330,7 @@ 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) + magics._make_bqstorage_client(True, credentials_mock, {}) error_msg = str(exc_context.value) assert "google-cloud-bigquery-storage" in error_msg @@ -347,7 +347,7 @@ def test__make_bqstorage_client_true_missing_gapic(missing_grpcio_lib): ) with pytest.raises(ImportError) as exc_context, missing_grpcio_lib: - magics._make_bqstorage_client(True, credentials_mock) + magics._make_bqstorage_client(True, credentials_mock, {}) assert "grpcio" in str(exc_context.value) @@ -1181,7 +1181,7 @@ def test_bigquery_magic_with_project(): @pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_with_api_endpoint(ipython_ns_cleanup): +def test_bigquery_magic_with_bigquery_api_endpoint(ipython_ns_cleanup): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._connection = None @@ -1191,34 +1191,83 @@ def test_bigquery_magic_with_api_endpoint(ipython_ns_cleanup): ) with run_query_patch as run_query_mock: ip.run_cell_magic( - "bigquery", "--api_endpoint=https://api.endpoint.com", "SELECT 17 as num" + "bigquery", + "--bigquery_api_endpoint=https://bigquery_api.endpoint.com", + "SELECT 17 as num", ) connection_used = run_query_mock.call_args_list[0][0][0]._connection - assert connection_used.API_BASE_URL == "https://api.endpoint.com" + assert connection_used.API_BASE_URL == "https://bigquery_api.endpoint.com" # context client options should not change - assert magics.context.client_options.api_endpoint is None + assert magics.context.bigquery_client_options.api_endpoint is None @pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_with_api_endpoint_context_dict(): +def test_bigquery_magic_with_bigquery_api_endpoint_context_dict(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._connection = None - magics.context.client_options = {} + magics.context.bigquery_client_options = {} run_query_patch = mock.patch( "google.cloud.bigquery.magics.magics._run_query", autospec=True ) with run_query_patch as run_query_mock: ip.run_cell_magic( - "bigquery", "--api_endpoint=https://api.endpoint.com", "SELECT 17 as num" + "bigquery", + "--bigquery_api_endpoint=https://bigquery_api.endpoint.com", + "SELECT 17 as num", ) connection_used = run_query_mock.call_args_list[0][0][0]._connection - assert connection_used.API_BASE_URL == "https://api.endpoint.com" + assert connection_used.API_BASE_URL == "https://bigquery_api.endpoint.com" + # context client options should not change + assert magics.context.bigquery_client_options == {} + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_with_storage_api_endpoint(ipython_ns_cleanup): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._connection = None + + run_query_patch = mock.patch( + "google.cloud.bigquery.magics.magics._run_query", autospec=True + ) + with run_query_patch as run_query_mock: + ip.run_cell_magic( + "bigquery", + "--storage_api_endpoint=https://storage_api.endpoint.com", + "SELECT 17 as num", + ) + + client_used = run_query_mock.mock_calls[1][2]["bqstorage_client"] + assert client_used._transport._host == "https://storage_api.endpoint.com" + # context client options should not change + assert magics.context.storage_client_options.api_endpoint is None + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_with_storage_api_endpoint_context_dict(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._connection = None + magics.context.storage_client_options = {} + + run_query_patch = mock.patch( + "google.cloud.bigquery.magics.magics._run_query", autospec=True + ) + with run_query_patch as run_query_mock: + ip.run_cell_magic( + "bigquery", + "--storage_api_endpoint=https://storage_api.endpoint.com", + "SELECT 17 as num", + ) + + client_used = run_query_mock.mock_calls[1][2]["bqstorage_client"] + assert client_used._transport._host == "https://storage_api.endpoint.com" # context client options should not change - assert magics.context.client_options == {} + assert magics.context.storage_client_options == {} @pytest.mark.usefixtures("ipython_interactive") From 97a32cfc96078864f36caba696a2c56d70c603f4 Mon Sep 17 00:00:00 2001 From: Carlos de la Guardia Date: Wed, 14 Oct 2020 13:22:13 -0500 Subject: [PATCH 3/3] use unambiguous name for bigquery storage client options --- google/cloud/bigquery/magics/magics.py | 32 +++++++++++++------------- tests/unit/test_magics.py | 18 +++++++-------- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index 1f2ee54c7..5645a84a5 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/google/cloud/bigquery/magics/magics.py @@ -181,7 +181,7 @@ def __init__(self): self._connection = None self._default_query_job_config = bigquery.QueryJobConfig() self._bigquery_client_options = client_options.ClientOptions() - self._storage_client_options = client_options.ClientOptions() + self._bqstorage_client_options = client_options.ClientOptions() @property def credentials(self): @@ -268,7 +268,7 @@ def bigquery_client_options(self, value): self._bigquery_client_options = value @property - def storage_client_options(self): + def bqstorage_client_options(self): """google.api_core.client_options.ClientOptions: client options to be used through IPython magics for the storage client. @@ -283,13 +283,13 @@ def storage_client_options(self): >>> from google.cloud.bigquery import magics >>> client_options = {} >>> client_options['api_endpoint'] = "https://some.special.url" - >>> magics.context.storage_client_options = client_options + >>> magics.context.bqstorage_client_options = client_options """ - return self._storage_client_options + return self._bqstorage_client_options - @storage_client_options.setter - def storage_client_options(self, value): - self._storage_client_options = value + @bqstorage_client_options.setter + def bqstorage_client_options(self, value): + self._bqstorage_client_options = value @property def default_query_job_config(self): @@ -472,12 +472,12 @@ def _create_dataset_if_necessary(client, dataset_id): ), ) @magic_arguments.argument( - "--storage_api_endpoint", + "--bqstorage_api_endpoint", type=str, default=None, help=( - "The desired API endpoint, e.g., bigquery.googlepis.com. Defaults to this " - "option's value in the context storage_client_options." + "The desired API endpoint, e.g., bigquerystorage.googlepis.com. Defaults to " + "this option's value in the context bqstorage_client_options." ), ) @magic_arguments.argument( @@ -599,15 +599,15 @@ def _cell_magic(line, query): if context._connection: client._connection = context._connection - storage_client_options = copy.deepcopy(context.storage_client_options) - if args.storage_api_endpoint: - if isinstance(storage_client_options, dict): - storage_client_options["api_endpoint"] = args.storage_api_endpoint + bqstorage_client_options = copy.deepcopy(context.bqstorage_client_options) + if args.bqstorage_api_endpoint: + if isinstance(bqstorage_client_options, dict): + bqstorage_client_options["api_endpoint"] = args.bqstorage_api_endpoint else: - storage_client_options.api_endpoint = args.storage_api_endpoint + bqstorage_client_options.api_endpoint = args.bqstorage_api_endpoint bqstorage_client = _make_bqstorage_client( - use_bqstorage_api, context.credentials, storage_client_options, + use_bqstorage_api, context.credentials, bqstorage_client_options, ) close_transports = functools.partial(_close_transports, client, bqstorage_client) diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py index a3eec636d..30ca4d70c 100644 --- a/tests/unit/test_magics.py +++ b/tests/unit/test_magics.py @@ -1226,7 +1226,7 @@ def test_bigquery_magic_with_bigquery_api_endpoint_context_dict(): @pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_with_storage_api_endpoint(ipython_ns_cleanup): +def test_bigquery_magic_with_bqstorage_api_endpoint(ipython_ns_cleanup): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._connection = None @@ -1237,22 +1237,22 @@ def test_bigquery_magic_with_storage_api_endpoint(ipython_ns_cleanup): with run_query_patch as run_query_mock: ip.run_cell_magic( "bigquery", - "--storage_api_endpoint=https://storage_api.endpoint.com", + "--bqstorage_api_endpoint=https://bqstorage_api.endpoint.com", "SELECT 17 as num", ) client_used = run_query_mock.mock_calls[1][2]["bqstorage_client"] - assert client_used._transport._host == "https://storage_api.endpoint.com" + assert client_used._transport._host == "https://bqstorage_api.endpoint.com" # context client options should not change - assert magics.context.storage_client_options.api_endpoint is None + assert magics.context.bqstorage_client_options.api_endpoint is None @pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_with_storage_api_endpoint_context_dict(): +def test_bigquery_magic_with_bqstorage_api_endpoint_context_dict(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._connection = None - magics.context.storage_client_options = {} + magics.context.bqstorage_client_options = {} run_query_patch = mock.patch( "google.cloud.bigquery.magics.magics._run_query", autospec=True @@ -1260,14 +1260,14 @@ def test_bigquery_magic_with_storage_api_endpoint_context_dict(): with run_query_patch as run_query_mock: ip.run_cell_magic( "bigquery", - "--storage_api_endpoint=https://storage_api.endpoint.com", + "--bqstorage_api_endpoint=https://bqstorage_api.endpoint.com", "SELECT 17 as num", ) client_used = run_query_mock.mock_calls[1][2]["bqstorage_client"] - assert client_used._transport._host == "https://storage_api.endpoint.com" + assert client_used._transport._host == "https://bqstorage_api.endpoint.com" # context client options should not change - assert magics.context.storage_client_options == {} + assert magics.context.bqstorage_client_options == {} @pytest.mark.usefixtures("ipython_interactive")