added max_results magic option and fixed broken tests

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)
 

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"]