feature: add query location for bigquery magic (#1771)

Co-authored-by: Lingqing Gan <[email protected]>

Author: Gaurang033

google/cloud/bigquery/magics/magics.py

@@ -508,6 +508,15 @@ def _create_dataset_if_necessary(client, dataset_id):
         "Defaults to use tqdm_notebook. Install the ``tqdm`` package to use this feature."
     ),
 )
+@magic_arguments.argument(
+    "--location",
+    type=str,
+    default=None,
+    help=(
+        "Set the location to execute query."
+        "Defaults to location set in query setting in console."
+    ),
+)
 def _cell_magic(line, query):
     """Underlying function for bigquery cell magic
 
@@ -551,6 +560,7 @@ def _cell_magic(line, query):
             category=DeprecationWarning,
         )
     use_bqstorage_api = not args.use_rest_api
+    location = args.location
 
     params = []
     if params_option_value:
@@ -579,6 +589,7 @@ def _cell_magic(line, query):
         default_query_job_config=context.default_query_job_config,
         client_info=client_info.ClientInfo(user_agent=IPYTHON_USER_AGENT),
         client_options=bigquery_client_options,
+        location=location,
     )
     if context._connection:
         client._connection = context._connection

tests/unit/test_magics.py

@@ -2053,3 +2053,21 @@ def test_bigquery_magic_create_dataset_fails():
         )
 
     assert close_transports.called
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_with_location():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context.credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    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", "--location=us-east1", "SELECT 17 AS num")
+
+        client_options_used = run_query_mock.call_args_list[0][0][0]
+        assert client_options_used.location == "us-east1"