feat: allow client options to be set in magics context

google/cloud/bigquery/magics/magics.py

@@ -180,7 +180,8 @@ def __init__(self):
         self._project = None
         self._connection = None
         self._default_query_job_config = bigquery.QueryJobConfig()
-        self._client_options = client_options.ClientOptions()
+        self._bigquery_client_options = client_options.ClientOptions()
+        self._storage_client_options = client_options.ClientOptions()
 
     @property
     def credentials(self):
@@ -243,28 +244,52 @@ def project(self, value):
         self._project = value
 
     @property
-    def client_options(self):
+    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://www.googleapis.com/ end point.
+            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.client_options = client_options
+            >>> magics.context.bigquery_client_options = client_options
         """
-        return self._client_options
+        return self._bigquery_client_options
 
-    @client_options.setter
-    def client_options(self, value):
-        self._client_options = value
+    @bigquery_client_options.setter
+    def bigquery_client_options(self, value):
+        self._bigquery_client_options = value
+
+    @property
+    def storage_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.storage_client_options = client_options
+        """
+        return self._storage_client_options
+
+    @storage_client_options.setter
+    def storage_client_options(self, value):
+        self._storage_client_options = value
 
     @property
     def default_query_job_config(self):
@@ -438,12 +463,21 @@ def _create_dataset_if_necessary(client, dataset_id):
     ),
 )
 @magic_arguments.argument(
-    "--api_endpoint",
+    "--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(
+    "--storage_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."
+        "The desired API endpoint, e.g., bigquery.googlepis.com. Defaults to this "
+        "option's value in the context storage_client_options."
     ),
 )
 @magic_arguments.argument(
@@ -548,23 +582,33 @@ def _cell_magic(line, query):
 
     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
+    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:
-            client_options.api_endpoint = args.api_endpoint
+            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=client_options,
+        client_options=bigquery_client_options,
     )
     if context._connection:
         client._connection = context._connection
-    bqstorage_client = _make_bqstorage_client(use_bqstorage_api, context.credentials)
+
+    storage_client_options = copy.deepcopy(context.storage_client_options)
+    if args.storage_api_endpoint:
+        if isinstance(storage_client_options, dict):
+            storage_client_options["api_endpoint"] = args.storage_api_endpoint
+        else:
+            storage_client_options.api_endpoint = args.storage_api_endpoint
+
+    bqstorage_client = _make_bqstorage_client(
+        use_bqstorage_api, context.credentials, storage_client_options,
+    )
 
     close_transports = functools.partial(_close_transports, client, bqstorage_client)
 
@@ -677,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
 
@@ -703,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)
 
@@ -1181,7 +1181,7 @@ def test_bigquery_magic_with_project():
 
 
 @pytest.mark.usefixtures("ipython_interactive")
-def test_bigquery_magic_with_api_endpoint(ipython_ns_cleanup):
+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
@@ -1191,34 +1191,83 @@ def test_bigquery_magic_with_api_endpoint(ipython_ns_cleanup):
     )
     with run_query_patch as run_query_mock:
         ip.run_cell_magic(
-            "bigquery", "--api_endpoint=https://api.endpoint.com", "SELECT 17 as num"
+            "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://api.endpoint.com"
+    assert connection_used.API_BASE_URL == "https://bigquery_api.endpoint.com"
     # context client options should not change
-    assert magics.context.client_options.api_endpoint is None
+    assert magics.context.bigquery_client_options.api_endpoint is None
 
 
 @pytest.mark.usefixtures("ipython_interactive")
-def test_bigquery_magic_with_api_endpoint_context_dict():
+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.client_options = {}
+    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", "--api_endpoint=https://api.endpoint.com", "SELECT 17 as num"
+            "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://api.endpoint.com"
+    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_storage_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",
+            "--storage_api_endpoint=https://storage_api.endpoint.com",
+            "SELECT 17 as num",
+        )
+
+    client_used = run_query_mock.mock_calls[1][2]["bqstorage_client"]
+    assert client_used._transport._host == "https://storage_api.endpoint.com"
+    # context client options should not change
+    assert magics.context.storage_client_options.api_endpoint is None
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_with_storage_api_endpoint_context_dict():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._connection = None
+    magics.context.storage_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",
+            "--storage_api_endpoint=https://storage_api.endpoint.com",
+            "SELECT 17 as num",
+        )
+
+    client_used = run_query_mock.mock_calls[1][2]["bqstorage_client"]
+    assert client_used._transport._host == "https://storage_api.endpoint.com"
     # context client options should not change
-    assert magics.context.client_options == {}
+    assert magics.context.storage_client_options == {}
 
 
 @pytest.mark.usefixtures("ipython_interactive")