cell magic accepts table_ids instead of queries

bigquery/google/cloud/bigquery/magics.py

@@ -129,6 +129,7 @@
 
 from __future__ import print_function
 
+import re
 import ast
 import sys
 import time
@@ -265,6 +266,13 @@ 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, max_results=None):
     """Runs a query while printing status updates
@@ -434,6 +442,26 @@ def _cell_magic(line, query):
     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
@@ -445,7 +473,6 @@ def _cell_magic(line, query):
         value = int(args.maximum_bytes_billed)
         job_config.maximum_bytes_billed = value
 
-    error = None
     try:
         query_job = _run_query(
             client, query, job_config=job_config, max_results=max_results
@@ -457,12 +484,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:

bigquery/tests/unit/test_magics.py

@@ -695,6 +695,92 @@ def test_bigquery_magic_w_max_results_valid_calls_queryjob_result():
 
         query_job_mock.result.assert_called_once_with(timeout=0.5, max_results=5)
 
+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", 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_dryrun_option_sets_job_config():