Skip to content
Permalink
Browse files
feat: allow client options to be set in magics context (#322)
* feat: allow client options to be set in magics context

* add separate client options for storage client
  • Loading branch information
cguardia committed Oct 14, 2020
1 parent fb401bd commit 5178b55682f5e264bfc082cde26acb1fdc953a18
Showing with 188 additions and 8 deletions.
  1. +94 −4 google/cloud/bigquery/magics/magics.py
  2. +94 −4 tests/unit/test_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,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
@@ -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
@@ -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
@@ -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",
@@ -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)

@@ -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

@@ -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,
)


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

@@ -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()

0 comments on commit 5178b55

Please sign in to comment.