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

Author: cguardia

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

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