feat: allow client options to be set in magics context

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

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