From 6559dde1568b2d07c1e0a142194f2b370efb8983 Mon Sep 17 00:00:00 2001 From: Gaurang Shah Date: Tue, 23 Jan 2024 19:03:36 -0500 Subject: [PATCH] feature: add query location for bigquery magic (#1771) Co-authored-by: Lingqing Gan --- google/cloud/bigquery/magics/magics.py | 11 +++++++++++ tests/unit/test_magics.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py index 2a3583c66..b7c685d9a 100644 --- a/google/cloud/bigquery/magics/magics.py +++ b/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 diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py index b03894095..1511cba9c 100644 --- a/tests/unit/test_magics.py +++ b/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"