feat: allow cell magic body to be a $variable (#1053)

* feat: allow cell magic body to be a $variable

* Fix missing indefinitive article in error msg

* Adjust test assertion to error message change

* Refactor logic for extracting query variable

* Explicitly warn about missing query variable name

* Thest the query "variable" is not identifier case

Author: plamut

google/cloud/bigquery/magics/magics.py

@@ -596,6 +596,29 @@ def _cell_magic(line, query):
             _handle_error(error, args.destination_var)
             return
 
+        # Check if query is given as a reference to a variable.
+        if query.startswith("$"):
+            query_var_name = query[1:]
+
+            if not query_var_name:
+                missing_msg = 'Missing query variable name, empty "$" is not allowed.'
+                raise NameError(missing_msg)
+
+            if query_var_name.isidentifier():
+                ip = IPython.get_ipython()
+                query = ip.user_ns.get(query_var_name, ip)  # ip serves as a sentinel
+
+                if query is ip:
+                    raise NameError(
+                        f"Unknown query, variable {query_var_name} does not exist."
+                    )
+                else:
+                    if not isinstance(query, (str, bytes)):
+                        raise TypeError(
+                            f"Query variable {query_var_name} must be a string "
+                            "or a bytes-like value."
+                        )
+
         # Any query that does not contain whitespace (aside from leading and trailing whitespace)
         # is assumed to be a table id
         if not re.search(r"\s", query):

tests/unit/test_magics.py

@@ -584,7 +584,7 @@ def test_bigquery_magic_does_not_clear_display_in_verbose_mode():
 
 
 @pytest.mark.usefixtures("ipython_interactive")
-def test_bigquery_magic_clears_display_in_verbose_mode():
+def test_bigquery_magic_clears_display_in_non_verbose_mode():
     ip = IPython.get_ipython()
     ip.extension_manager.load_extension("google.cloud.bigquery")
     magics.context.credentials = mock.create_autospec(
@@ -1710,6 +1710,143 @@ def test_bigquery_magic_with_improperly_formatted_params():
         ip.run_cell_magic("bigquery", "--params {17}", sql)
 
 
+@pytest.mark.parametrize(
+    "raw_sql", ("SELECT answer AS 42", " \t   SELECT answer AS 42  \t  ")
+)
+@pytest.mark.usefixtures("ipython_interactive")
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_bigquery_magic_valid_query_in_existing_variable(ipython_ns_cleanup, raw_sql):
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context.credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    ipython_ns_cleanup.append((ip, "custom_query"))
+    ipython_ns_cleanup.append((ip, "query_results_df"))
+
+    run_query_patch = mock.patch(
+        "google.cloud.bigquery.magics.magics._run_query", autospec=True
+    )
+    query_job_mock = mock.create_autospec(
+        google.cloud.bigquery.job.QueryJob, instance=True
+    )
+    mock_result = pandas.DataFrame([42], columns=["answer"])
+    query_job_mock.to_dataframe.return_value = mock_result
+
+    ip.user_ns["custom_query"] = raw_sql
+    cell_body = "$custom_query"  # Referring to an existing variable name (custom_query)
+    assert "query_results_df" not in ip.user_ns
+
+    with run_query_patch as run_query_mock:
+        run_query_mock.return_value = query_job_mock
+
+        ip.run_cell_magic("bigquery", "query_results_df", cell_body)
+
+        run_query_mock.assert_called_once_with(mock.ANY, raw_sql, mock.ANY)
+
+    assert "query_results_df" in ip.user_ns  # verify that the variable exists
+    df = ip.user_ns["query_results_df"]
+    assert len(df) == len(mock_result)  # verify row count
+    assert list(df) == list(mock_result)  # verify column names
+    assert list(df["answer"]) == [42]
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_bigquery_magic_nonexisting_query_variable():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context.credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    run_query_patch = mock.patch(
+        "google.cloud.bigquery.magics.magics._run_query", autospec=True
+    )
+
+    ip.user_ns.pop("custom_query", None)  # Make sure the variable does NOT exist.
+    cell_body = "$custom_query"  # Referring to a non-existing variable name.
+
+    with pytest.raises(
+        NameError, match=r".*custom_query does not exist.*"
+    ), run_query_patch as run_query_mock:
+        ip.run_cell_magic("bigquery", "", cell_body)
+
+    run_query_mock.assert_not_called()
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_bigquery_magic_empty_query_variable_name():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context.credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    run_query_patch = mock.patch(
+        "google.cloud.bigquery.magics.magics._run_query", autospec=True
+    )
+    cell_body = "$"  # Not referring to any variable (name omitted).
+
+    with pytest.raises(
+        NameError, match=r"(?i).*missing query variable name.*"
+    ), run_query_patch as run_query_mock:
+        ip.run_cell_magic("bigquery", "", cell_body)
+
+    run_query_mock.assert_not_called()
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_bigquery_magic_query_variable_non_string(ipython_ns_cleanup):
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context.credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    run_query_patch = mock.patch(
+        "google.cloud.bigquery.magics.magics._run_query", autospec=True
+    )
+
+    ipython_ns_cleanup.append((ip, "custom_query"))
+
+    ip.user_ns["custom_query"] = object()
+    cell_body = "$custom_query"  # Referring to a non-string variable.
+
+    with pytest.raises(
+        TypeError, match=r".*must be a string or a bytes-like.*"
+    ), run_query_patch as run_query_mock:
+        ip.run_cell_magic("bigquery", "", cell_body)
+
+    run_query_mock.assert_not_called()
+
+
+@pytest.mark.usefixtures("ipython_interactive")
+@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
+def test_bigquery_magic_query_variable_not_identifier():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("google.cloud.bigquery")
+    magics.context.credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    cell_body = "$123foo"  # 123foo is not valid Python identifier
+
+    with io.capture_output() as captured_io:
+        ip.run_cell_magic("bigquery", "", cell_body)
+
+    # If "$" prefixes a string that is not a Python identifier, we do not treat such
+    # cell_body as a variable reference and just treat is as any other cell body input.
+    # If at the same time the cell body does not contain any whitespace, it is
+    # considered a table name, thus we expect an error that the table ID is not valid.
+    output = captured_io.stderr
+    assert "ERROR:" in output
+    assert "must be a fully-qualified ID" in output
+
+
 @pytest.mark.usefixtures("ipython_interactive")
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
 def test_bigquery_magic_with_invalid_multiple_option_values():