removed max_results param from run_query, updated tests

googleapis · Sep 4, 2019 · f88ffe8 · f88ffe8
1 parent 37b4d39
commit f88ffe8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 21 deletions.
diff --git 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
@@ -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"]