Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: allow client options to be set in magics context #322

Merged
merged 8 commits into from
Oct 14, 2020
98 changes: 94 additions & 4 deletions google/cloud/bigquery/magics/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@

import re
import ast
import copy
import functools
import sys
import time
Expand All @@ -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
Expand All @@ -178,11 +180,13 @@ def __init__(self):
self._project = None
self._connection = None
self._default_query_job_config = bigquery.QueryJobConfig()
self._bigquery_client_options = client_options.ClientOptions()
self._bqstorage_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
Expand Down Expand Up @@ -217,7 +221,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
Expand All @@ -239,6 +243,54 @@ def project(self):
def project(self, value):
self._project = value

@property
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://bigquery.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.bigquery_client_options = client_options
"""
return self._bigquery_client_options

@bigquery_client_options.setter
def bigquery_client_options(self, value):
self._bigquery_client_options = value

@property
def bqstorage_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.bqstorage_client_options = client_options
"""
return self._bqstorage_client_options

@bqstorage_client_options.setter
def bqstorage_client_options(self, value):
self._bqstorage_client_options = value

@property
def default_query_job_config(self):
"""google.cloud.bigquery.job.QueryJobConfig: Default job
Expand Down Expand Up @@ -410,6 +462,24 @@ def _create_dataset_if_necessary(client, dataset_id):
"Standard SQL if this argument is not used."
),
)
@magic_arguments.argument(
"--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(
"--bqstorage_api_endpoint",
type=str,
default=None,
help=(
"The desired API endpoint, e.g., bigquerystorage.googlepis.com. Defaults to "
"this option's value in the context bqstorage_client_options."
),
)
@magic_arguments.argument(
"--use_bqstorage_api",
action="store_true",
Expand Down Expand Up @@ -511,15 +581,34 @@ def _cell_magic(line, query):
params = _helpers.to_query_parameters(ast.literal_eval(params_option_value))

project = args.project or context.project

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:
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=bigquery_client_options,
)
if context._connection:
client._connection = context._connection
bqstorage_client = _make_bqstorage_client(use_bqstorage_api, context.credentials)

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:
bqstorage_client_options.api_endpoint = args.bqstorage_api_endpoint

bqstorage_client = _make_bqstorage_client(
use_bqstorage_api, context.credentials, bqstorage_client_options,
)

close_transports = functools.partial(_close_transports, client, bqstorage_client)

Expand Down Expand Up @@ -632,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

Expand All @@ -658,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,
)


Expand Down
98 changes: 94 additions & 4 deletions tests/unit/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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)


Expand All @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -1180,6 +1180,96 @@ def test_bigquery_magic_with_project():
assert magics.context.project == "general-project"


@pytest.mark.usefixtures("ipython_interactive")
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

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",
"--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://bigquery_api.endpoint.com"
# context client options should not change
assert magics.context.bigquery_client_options.api_endpoint is None


@pytest.mark.usefixtures("ipython_interactive")
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.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",
"--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://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_bqstorage_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",
"--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://bqstorage_api.endpoint.com"
# context client options should not change
assert magics.context.bqstorage_client_options.api_endpoint is None


@pytest.mark.usefixtures("ipython_interactive")
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.bqstorage_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",
"--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://bqstorage_api.endpoint.com"
# context client options should not change
assert magics.context.bqstorage_client_options == {}


@pytest.mark.usefixtures("ipython_interactive")
def test_bigquery_magic_with_multiple_options():
ip = IPython.get_ipython()
Expand Down