From 81b0ca2c43d68fa61158f5c3bd37bd703af17fb8 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Tue, 3 Sep 2019 17:53:48 -0700 Subject: [PATCH 1/6] added max_results magic option and fixed broken tests --- bigquery/google/cloud/bigquery/magics.py | 22 +++++++++++++++++++--- bigquery/tests/unit/test_magics.py | 8 ++++---- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index b2dae2511ec8..1d6a1c9935b0 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -266,7 +266,7 @@ def default_query_job_config(self, value): context = Context() -def _run_query(client, query, job_config=None): +def _run_query(client, query, job_config=None, max_results=None): """Runs a query while printing status updates Args: @@ -300,7 +300,7 @@ def _run_query(client, query, job_config=None): while True: print("\rQuery executing: {:0.2f}s".format(time.time() - start_time), end="") try: - query_job.result(timeout=0.5) + query_job.result(timeout=0.5, max_results=max_results) break except futures.TimeoutError: continue @@ -320,6 +320,16 @@ def _run_query(client, query, job_config=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( "--maximum_bytes_billed", default=None, @@ -420,6 +430,12 @@ def _cell_magic(line, query): bqstorage_client = _make_bqstorage_client( args.use_bqstorage_api or context.use_bqstorage_api, context.credentials ) + + if args.max_results: + max_results = int(args.max_results) + else: + max_results = None + job_config = bigquery.job.QueryJobConfig() job_config.query_parameters = params job_config.use_legacy_sql = args.use_legacy_sql @@ -433,7 +449,7 @@ def _cell_magic(line, query): error = None try: - query_job = _run_query(client, query, job_config) + 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 fbea9bdd9050..fac0fee81600 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][0][-1] + job_config_used = run_query_mock.call_args_list[0][1]['job_config'] assert job_config_used.use_legacy_sql is True @@ -662,7 +662,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][0][-1] + job_config_used = run_query_mock.call_args_list[0][1]['job_config'] assert job_config_used.dry_run is True @@ -924,7 +924,7 @@ 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) + 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 +959,7 @@ 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) + 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"] From b04cc374900957ec2b0a8477be35cdadb67fb854 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Tue, 3 Sep 2019 18:36:23 -0700 Subject: [PATCH 2/6] added tests for --max_results magic option --- bigquery/google/cloud/bigquery/magics.py | 6 +-- bigquery/tests/unit/test_magics.py | 63 ++++++++++++++++++++++-- 2 files changed, 62 insertions(+), 7 deletions(-) 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"] From 7176592cd82b961c75b2b0f3c7abfbb4d1545686 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Tue, 3 Sep 2019 17:53:48 -0700 Subject: [PATCH 3/6] 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"] From 37b4d393d832187a42034ba4b7f4516849901aba Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Tue, 3 Sep 2019 18:36:23 -0700 Subject: [PATCH 4/6] added tests for --max_results magic option --- bigquery/google/cloud/bigquery/magics.py | 2 -- bigquery/tests/unit/test_magics.py | 14 ++------------ 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 67583150ab4c..91ff2b56b7a2 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( "--max_results", default=None, diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 5c7937125e4d..5e0745e2e33f 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -414,11 +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") -<<<<<<< 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 @@ -979,13 +975,10 @@ 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"] @@ -1020,13 +1013,10 @@ 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"] From 147e43d54c907a8fc6ee7e06d56aaf7f2a4a5a47 Mon Sep 17 00:00:00 2001 From: Shubha Rajan Date: Wed, 4 Sep 2019 13:43:03 -0700 Subject: [PATCH 5/6] Removed duplicate `--max_results` magic argument --- bigquery/google/cloud/bigquery/magics.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 91ff2b56b7a2..ab5fb4149731 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -328,14 +328,6 @@ def _run_query(client, query, job_config=None, max_results=None): "Defaults to returning all rows." ), ) -@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( "--maximum_bytes_billed", default=None, From f88ffe86a106da05c31939736d51a77f1fe50f52 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Wed, 4 Sep 2019 14:29:12 -0700 Subject: [PATCH 6/6] removed max_results param from run_query, updated tests --- bigquery/google/cloud/bigquery/magics.py | 24 ++++++++++-------------- bigquery/tests/unit/test_magics.py | 10 +++------- 2 files changed, 13 insertions(+), 21 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 91ff2b56b7a2..4c93d1307a42 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -266,7 +266,7 @@ def default_query_job_config(self, value): context = Context() -def _run_query(client, query, job_config=None, max_results=None): +def _run_query(client, query, job_config=None): """Runs a query while printing status updates Args: @@ -300,7 +300,7 @@ def _run_query(client, query, job_config=None, max_results=None): while True: print("\rQuery executing: {:0.2f}s".format(time.time() - start_time), end="") try: - query_job.result(timeout=0.5, max_results=max_results) + query_job.result(timeout=0.5) break except futures.TimeoutError: continue @@ -328,14 +328,6 @@ def _run_query(client, query, job_config=None, max_results=None): "Defaults to returning all rows." ), ) -@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( "--maximum_bytes_billed", default=None, @@ -455,9 +447,7 @@ 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) except Exception as ex: error = str(ex) @@ -484,7 +474,13 @@ def _cell_magic(line, query): ) return query_job - result = query_job.to_dataframe(bqstorage_client=bqstorage_client) + if max_results: + result = query_job.result(max_results=max_results).to_dataframe( + bqstorage_client=bqstorage_client + ) + else: + result = query_job.to_dataframe(bqstorage_client=bqstorage_client) + if args.destination_var: IPython.get_ipython().push({args.destination_var: result}) else: diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 5e0745e2e33f..ed748d2dd5e3 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -693,7 +693,7 @@ def test_bigquery_magic_w_max_results_valid_calls_queryjob_result(): 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) + query_job_mock.result.assert_called_with(max_results=5) @pytest.mark.usefixtures("ipython_interactive") @@ -976,9 +976,7 @@ def test_bigquery_magic_with_string_params(): 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) assert "params_string_df" in ip.user_ns # verify that the variable exists df = ip.user_ns["params_string_df"] @@ -1014,9 +1012,7 @@ def test_bigquery_magic_with_dict_params(): 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) assert "params_dict_df" in ip.user_ns # verify that the variable exists df = ip.user_ns["params_dict_df"]