From 7176592cd82b961c75b2b0f3c7abfbb4d1545686 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Tue, 3 Sep 2019 17:53:48 -0700 Subject: [PATCH] added max_results magic option and fixed broken tests --- bigquery/google/cloud/bigquery/magics.py | 10 ++++++++++ bigquery/tests/unit/test_magics.py | 12 ++++++++++++ 2 files changed, 22 insertions(+) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index ab5fb4149731..67583150ab4c 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -320,6 +320,16 @@ def _run_query(client, query, job_config=None, max_results=None): default=None, help=("Project to use for executing this query. Defaults to the context project."), ) + +@magic_arguments.argument( + "--max_results", + default=None, + help=( + "Maximum number of rows in dataframe returned from executing the query." + "Defaults to returning all rows." + ), +) + @magic_arguments.argument( "--max_results", default=None, diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index e70775db571b..5c7937125e4d 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -414,7 +414,11 @@ def test_bigquery_magic_with_legacy_sql(): with run_query_patch as run_query_mock: ip.run_cell_magic("bigquery", "--use_legacy_sql", "SELECT 17 AS num") +<<<<<<< HEAD job_config_used = run_query_mock.call_args_list[0][1]["job_config"] +======= + job_config_used = run_query_mock.call_args_list[0][1]['job_config'] +>>>>>>> d4fa587dacf... added max_results magic option and fixed broken tests assert job_config_used.use_legacy_sql is True @@ -975,9 +979,13 @@ def test_bigquery_magic_with_string_params(): run_query_mock.return_value = query_job_mock ip.run_cell_magic("bigquery", 'params_string_df --params {"num":17}', sql) +<<<<<<< HEAD run_query_mock.assert_called_once_with( mock.ANY, sql.format(num=17), mock.ANY, max_results=None ) +======= + run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY, max_results=None) +>>>>>>> d4fa587dacf... added max_results magic option and fixed broken tests assert "params_string_df" in ip.user_ns # verify that the variable exists df = ip.user_ns["params_string_df"] @@ -1012,9 +1020,13 @@ def test_bigquery_magic_with_dict_params(): # Insert dictionary into user namespace so that it can be expanded ip.user_ns["params"] = params ip.run_cell_magic("bigquery", "params_dict_df --params $params", sql) +<<<<<<< HEAD run_query_mock.assert_called_once_with( mock.ANY, sql.format(num=17), mock.ANY, max_results=None ) +======= + run_query_mock.assert_called_once_with(mock.ANY, sql.format(num=17), mock.ANY, max_results=None) +>>>>>>> d4fa587dacf... added max_results magic option and fixed broken tests assert "params_dict_df" in ip.user_ns # verify that the variable exists df = ip.user_ns["params_dict_df"]