Repository: googleapis/python-bigquery

Issue:
The magics API do not support setting client options

#### Environment details

  - OS type and version: jupyter notebook
  - Python version: Python 3.7.3
  - pip version: pip 20.1.1
  - google-cloud-bigquery version: 1.25.0

#### Steps to reproduce
This is a missing option that do exists in google bigquery client - "https://googleapis.dev/python/bigquery/latest/generated/google.cloud.bigquery.client.Client.html#google.cloud.bigquery.client.Client"
But is missing from "https://github.com/googleapis/python-bigquery/blob/master/google/cloud/bigquery/magics.py"

#### Code example

```python
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),
    )
```
Note the the client_option is missing in creating the client in the code above

The patch to apply to fix this flow:
```diff
diff --git a/google/cloud/bigquery/magics.py b/google/cloud/bigquery/magics.py
index 7128e32..00e26d0 100644
--- a/google/cloud/bigquery/magics.py
+++ b/google/cloud/bigquery/magics.py
@@ -140,6 +140,10 @@
         Out[7]:     num
            ...: -------
            ...: 0    17
+        Executing query with job ID: bf633912-af2c-4780-b568-5d868058632b
+        Query executing: 2.61s
+        Query complete after 2.92s
+
 """
 
 from __future__ import print_function
@@ -183,6 +187,7 @@ class Context(object):
         self._project = None
         self._connection = None
         self._default_query_job_config = bigquery.QueryJobConfig()
+        self._client_options = google.api_core.client_options.ClientOptions()
 
     @property
     def credentials(self):
@@ -244,6 +249,30 @@ class Context(object):
     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 does not need to be explicitly defined if you have an
+            no special network conenctions are required. Normally you would
+            be using https://www.googleapis.com/ end point
+
+        Example:
+            Manually setting the endpoint:
+
+            >>> from google.cloud.bigquery import magics
+            >>> client_option = {}
+            >>> client_option['api_endpoint'] = "https://some.special.url"
+            >>> magics.context.client_options = client_option
+        """
+        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
@@ -459,6 +488,18 @@ def _create_dataset_if_necessary(client, dataset_id):
         "name (ex. $my_dict_var)."
     ),
 )
+@magic_arguments.argument(
+    "--options",
+    nargs="+",
+    default=None,
+    help=(
+        "Parameters to format the query string. If present, the --options "
+        "flag should be followed by a string representation of a dictionary "
+        "in the format (Union[google.api_core.client_options.ClientOptions, Dict])"
+        "or a reference to a dictionary in the same format. The dictionary "
+        "key/value matches does at google.api_core.client_options.ClientOptions"
+    ),
+)
 def _cell_magic(line, query):
     """Underlying function for bigquery cell magic
 
@@ -501,6 +542,7 @@ def _cell_magic(line, query):
         credentials=context.credentials,
         default_query_job_config=context.default_query_job_config,
         client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
+        client_options=context.client_options
     )
     if context._connection:
         client._connection = context._connection
```


Discussion:
@boazsade Thanks for the suggestion and the patch!

I'll just check if Product has any objections (cc: @shollyman), and if not, we'll add this.

P.S.: If confirmed, would you mind opening an actual PR with added tests, or are you fine with us taking the patch as a base and create a PR ourselves?
(if significant portions of the patch are used, it would probably still require signing the CLA, though)
You can take this pach, this is something that we really need to be part of the code base, thanks for the reply 
@plamut  is there any progress with this issue? I was unable to create pull request, what about applying the patch? was that approved? thanks
@boazsade Not at this point, unfortunately, but I can re-check at this week's meeting. Lately one of the primary things in focus are preparations to dropping Python 2 support and moving them to the new code generator (for those who have some auto-generated parts).
Hi @plamut. I wanted to check in and see if there's an update regarding including this patch in an upcoming release. One of our customers is in the process of migrating their data environment to BigQuery and this has become a priority. Thanks in advance.
@yoavcloud Hi, I'm working on a PR to add your patch. Should be ready for review in a couple of days.
Thank you so much @cguardia!
@yoavcloud Well, I've been looking at this and have a couple of questions. While `api_endpoint` does seem like a useful option to have here, I don't think the rest of the client options are practical for magics, with the possible exception of `quota_project_id`. Passing in callables in the command line is just not possible, plus getting the lexer to parse the dictionary format in the first place would not be trivial. Have you tried the patch? I don't believe it can work as it is. We could probably add the api endpoint as an option, though. Would that cover your use case?
The patch file was generated by me (@boazsade), this patch was generate after we tested this to work. This was the reason we wanted it - after verify it. What we really need is the option to set the end point value (so that the traffic can pass not directly to google but through other authenticated host. As for the other parameters (like quota_project_id) I didn't test that, the reason I wrote it the way I did was to keep it with the same as the internal calls, but this can be dropped - again we really need the option to set the api_endpoint, other parameters as less important for us.
> I don't think the rest of the client options are practical for magics

We don't necessarily need command-line options for all of these (though it wouldn't hurt to add). The most important thing is to make sure it is settable from `google.cloud.bigquery.magics.context`, as this object sets default values for all calls to the magics.

Resolving patch:
diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py
--- 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,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,
     )
 
 


Test patch:
diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py
--- 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)
 
@@ -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()