shubha-rajan · Aug 28, 2019 · Aug 29, 2019 · Aug 29, 2019 · Aug 30, 2019 · Aug 30, 2019
diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py
@@ -1070,7 +1070,6 @@ def _get_query_results(
                 (Optional) number of milliseconds the the API call should
                 wait for the query to complete before the request times out.
             location (str): Location of the query job.
-
         Returns:
             google.cloud.bigquery.query._QueryResults:
                 A new ``_QueryResults`` instance.

diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -499,8 +499,8 @@ def _copy_configuration_properties(self, configuration):
     def _set_properties(self, api_response):
         """Update properties from resource in body of ``api_response``
 
-        :type api_response: dict
-        :param api_response: response returned from an API call
+        Args:
+            api_response (Dict): response returned from an API call.
         """
         cleaned = api_response.copy()
         self._scrub_local_properties(cleaned)
@@ -2151,6 +2151,19 @@ def flatten_results(self):
     def flatten_results(self, value):
         self._set_sub_prop("flattenResults", value)
 
+    @property
+    def max_results(self):
+        """int: Maximum number of results to read.
+
+        See
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query
+        """
+        return _helpers._int_or_none(self._get_sub_prop("maxResults"))
+
+    @max_results.setter
+    def max_results(self, value):
+        self._set_sub_prop("maxResults", value)
+
     @property
     def maximum_billing_tier(self):
         """int: Deprecated. Changes the billing tier to allow high-compute
@@ -2385,6 +2398,7 @@ def __init__(self, job_id, query, client, job_config=None):
 
         if job_config is None:
             job_config = QueryJobConfig()
+
         if job_config.use_legacy_sql is None:
             job_config.use_legacy_sql = False
 
@@ -2451,6 +2465,13 @@ def flatten_results(self):
         """
         return self._configuration.flatten_results
 
+    @property
+    def max_results(self):
+        """See
+        :attr:`google.cloud.bigquery.job.QueryJobConfig.max_results`.
+        """
+        return self._configuration.max_results
+
     @property
     def priority(self):
         """See
@@ -2558,6 +2579,18 @@ def to_api_repr(self):
 
         return resource
 
+    def _set_properties(self, api_response):
+        """Update properties from resource in body of ``api_response``
+
+        Args:
+            api_response (Dict): response returned from an API call.
+        """
+
+        # prevent max_results value from being erased
+        max_results = self.max_results
+        super(QueryJob, self)._set_properties(api_response)
+        self._configuration.max_results = max_results
+
     def _copy_configuration_properties(self, configuration):
         """Helper:  assign subclass configuration properties in cleaned."""
         self._configuration._properties = copy.deepcopy(configuration)
@@ -2933,6 +2966,7 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY):
             concurrent.futures.TimeoutError:
                 If the job did not complete in the given timeout.
         """
+
         try:
             super(QueryJob, self).result(timeout=timeout)
 
@@ -2956,7 +2990,9 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY):
         dest_table_ref = self.destination
         dest_table = Table(dest_table_ref, schema=schema)
         dest_table._properties["numRows"] = self._query_results.total_rows
-        rows = self._client.list_rows(dest_table, page_size=page_size, retry=retry)
+        rows = self._client.list_rows(
+            dest_table, page_size=page_size, retry=retry, max_results=self.max_results
+        )
         rows._preserve_order = _contains_order_by(self.query)
         return rows
 

diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py
@@ -129,6 +129,7 @@
 
 from __future__ import print_function
 
+import re
 import ast
 import sys
 import time
@@ -266,6 +267,15 @@ def default_query_job_config(self, value):
 context = Context()
 
 
+def _print_error(error, destination_var=None):
+    if destination_var:
+        print(
+            "Could not save output to variable '{}'.".format(destination_var),
+            file=sys.stderr,
+        )
+    print("\nERROR:\n", error, file=sys.stderr)
+
+
 def _run_query(client, query, job_config=None):
     """Runs a query while printing status updates
 
@@ -278,6 +288,7 @@ def _run_query(client, query, job_config=None):
         job_config (google.cloud.bigquery.job.QueryJobConfig, optional):
             Extra configuration options for the job.
 
+
     Returns:
         google.cloud.bigquery.job.QueryJob: the query job created
 
@@ -320,6 +331,14 @@ 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,
@@ -394,6 +413,7 @@ def _cell_magic(line, query):
     Returns:
         pandas.DataFrame: the query results.
     """
+
     args = magic_arguments.parse_argstring(_cell_magic, line)
 
     params = []
@@ -409,6 +429,7 @@ def _cell_magic(line, query):
             )
 
     project = args.project or context.project
+
     client = bigquery.Client(
         project=project,
         credentials=context.credentials,
@@ -420,18 +441,44 @@ 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
+
+    error = None
+
+    if not re.search(r"\s", query.rstrip()):
+        table_id = query.rstrip()
+
+        try:
+            rows = client.list_rows(table_id, max_results=max_results)
+        except Exception as ex:
+            error = str(ex)
+        if error:
+            _print_error(error, args.destination_var)
+            return
+
+        result = rows.to_dataframe(bqstorage_client=bqstorage_client)
+        if args.destination_var:
+            IPython.get_ipython().push({args.destination_var: result})
+            return
+        else:
+            return result
+
     job_config = bigquery.job.QueryJobConfig()
     job_config.query_parameters = params
     job_config.use_legacy_sql = args.use_legacy_sql
     job_config.dry_run = args.dry_run
+    job_config.max_results = max_results
 
     if args.maximum_bytes_billed == "None":
         job_config.maximum_bytes_billed = 0
     elif args.maximum_bytes_billed is not None:
         value = int(args.maximum_bytes_billed)
         job_config.maximum_bytes_billed = value
 
-    error = None
     try:
         query_job = _run_query(client, query, job_config)
     except Exception as ex:
@@ -441,12 +488,7 @@ def _cell_magic(line, query):
         display.clear_output()
 
     if error:
-        if args.destination_var:
-            print(
-                "Could not save output to variable '{}'.".format(args.destination_var),
-                file=sys.stderr,
-            )
-        print("\nERROR:\n", error, file=sys.stderr)
+        _print_error(error, args.destination_var)
         return
 
     if args.dry_run and args.destination_var:

diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py
@@ -743,6 +743,168 @@ def test_bigquery_magic_dryrun_option_saves_query_job_to_variable():
     assert isinstance(q_job, job.QueryJob)
 
 
+@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_sets_job_config():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    run_query_patch = mock.patch(
+        "google.cloud.bigquery.magics._run_query", autospec=True
+    )
+
+    sql = "SELECT 17 AS num"
+
+    with run_query_patch as run_query_mock:
+        ip.run_cell_magic("bigquery", "--max_results=5", sql)
+
+        job_config_used = run_query_mock.call_args_list[0][0][-1]
+        assert job_config_used.max_results == 5
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_w_table_id_and_max_results():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    row_iterator_mock = mock.create_autospec(
+        google.cloud.bigquery.table.RowIterator, instance=True
+    )
+
+    client_patch = mock.patch(
+        "google.cloud.bigquery.magics.bigquery.Client", autospec=True
+    )
+
+    run_query_patch = mock.patch(
+        "google.cloud.bigquery.magics._run_query", autospec=True
+    )
+
+    table_id = "bigquery-public-data.samples.shakespeare"
+    result = pandas.DataFrame([17], columns=["num"])
+
+    with client_patch as client_mock, run_query_patch as run_query_mock:
+        client_mock().list_rows.return_value = row_iterator_mock
+        row_iterator_mock.to_dataframe.return_value = result
+
+        return_value = ip.run_cell_magic("bigquery", "--max_results=5", table_id)
+        client_mock.return_value.list_rows.assert_called_once_with(
+            table_id, max_results=5
+        )
+        run_query_mock.assert_not_called()
+
+    assert isinstance(return_value, pandas.DataFrame)
+
+
+def test_bigquery_magic_w_table_id_invalid():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    list_rows_patch = mock.patch(
+        "google.cloud.bigquery.magics.bigquery.Client.list_rows",
+        autospec=True,
+        side_effect=exceptions.BadRequest("Not a valid table ID"),
+    )
+
+    table_id = "not-a-real-table"
+
+    with list_rows_patch, io.capture_output() as captured_io:
+        ip.run_cell_magic("bigquery", "df", table_id)
+
+    output = captured_io.stderr
+    assert "Could not save output to variable" in output
+    assert "400 Not a valid table ID" in output
+    assert "Traceback (most recent call last)" not in output
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_w_table_id_and_destination_var():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    row_iterator_mock = mock.create_autospec(
+        google.cloud.bigquery.table.RowIterator, instance=True
+    )
+
+    client_patch = mock.patch(
+        "google.cloud.bigquery.magics.bigquery.Client", autospec=True
+    )
+
+    table_id = "bigquery-public-data.samples.shakespeare"
+    result = pandas.DataFrame([17], columns=["num"])
+
+    with client_patch as client_mock:
+        client_mock().list_rows.return_value = row_iterator_mock
+        row_iterator_mock.to_dataframe.return_value = result
+
+        ip.run_cell_magic("bigquery", "df --max_results=5", table_id)
+
+    assert "df" in ip.user_ns
+    df = ip.user_ns["df"]
+
+    assert isinstance(df, pandas.DataFrame)
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_w_table_id_and_bqstorage_client():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context._project = None
+
+    row_iterator_mock = mock.create_autospec(
+        google.cloud.bigquery.table.RowIterator, instance=True
+    )
+
+    client_patch = mock.patch(
+        "google.cloud.bigquery.magics.bigquery.Client", autospec=True
+    )
+
+    bqstorage_mock = mock.create_autospec(
+        bigquery_storage_v1beta1.BigQueryStorageClient
+    )
+    bqstorage_instance_mock = mock.create_autospec(
+        bigquery_storage_v1beta1.BigQueryStorageClient, instance=True
+    )
+    bqstorage_mock.return_value = bqstorage_instance_mock
+    bqstorage_client_patch = mock.patch(
+        "google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient", bqstorage_mock
+    )
+
+    table_id = "bigquery-public-data.samples.shakespeare"
+
+    with client_patch as client_mock, bqstorage_client_patch:
+        client_mock().list_rows.return_value = row_iterator_mock
+
+        ip.run_cell_magic("bigquery", "--use_bqstorage_api --max_results=5", table_id)
+        row_iterator_mock.to_dataframe.assert_called_once_with(
+            bqstorage_client=bqstorage_instance_mock
+        )
+
+
 @pytest.mark.usefixtures("ipython_interactive")
 def test_bigquery_magic_w_maximum_bytes_billed_invalid():
     ip = IPython.get_ipython()