diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 1d6a1c9935b0..ab5fb4149731 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -320,7 +320,6 @@ 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, @@ -329,7 +328,6 @@ def _run_query(client, query, job_config=None, max_results=None): "Defaults to returning all rows." ), ) - @magic_arguments.argument( "--maximum_bytes_billed", default=None, @@ -449,7 +447,9 @@ def _cell_magic(line, query): error = None try: - query_job = _run_query(client, query, job_config=job_config, max_results=max_results) + query_job = _run_query( + client, query, job_config=job_config, max_results=max_results + ) except Exception as ex: error = str(ex) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index fac0fee81600..e70775db571b 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -414,7 +414,7 @@ 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") - 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"] assert job_config_used.use_legacy_sql is True @@ -645,6 +645,57 @@ def test_bigquery_magic_without_bqstorage(monkeypatch): assert isinstance(return_value, pandas.DataFrame) +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_max_results_invalid(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + client_query_patch = mock.patch( + "google.cloud.bigquery.client.Client.query", autospec=True + ) + + sql = "SELECT 17 AS num" + + with pytest.raises(ValueError), default_patch, client_query_patch: + ip.run_cell_magic("bigquery", "--max_results=abc", sql) + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_max_results_valid_calls_queryjob_result(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + client_query_patch = mock.patch( + "google.cloud.bigquery.client.Client.query", autospec=True + ) + + sql = "SELECT 17 AS num" + + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + + with client_query_patch as client_query_mock, default_patch: + client_query_mock.return_value = query_job_mock + ip.run_cell_magic("bigquery", "--max_results=5", sql) + + query_job_mock.result.assert_called_once_with(timeout=0.5, max_results=5) + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_dryrun_option_sets_job_config(): ip = IPython.get_ipython() @@ -662,7 +713,7 @@ def test_bigquery_magic_dryrun_option_sets_job_config(): with run_query_patch as run_query_mock: ip.run_cell_magic("bigquery", "--dry_run", sql) - 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"] assert job_config_used.dry_run is True @@ -924,7 +975,9 @@ 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) - 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 + ) assert "params_string_df" in ip.user_ns # verify that the variable exists df = ip.user_ns["params_string_df"] @@ -959,7 +1012,9 @@ 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) - 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 + ) assert "params_dict_df" in ip.user_ns # verify that the variable exists df = ip.user_ns["params_dict_df"]