Repository: googleapis/python-bigquery

Issue:
allow the body of `%%bigquery` cell magic to be a `$variable`
 **Is your feature request related to a problem? Please describe.**

When using BigQuery notebooks. There are queries that would be nice to be parametrized, but the parameters are in places where BigQuery parameterized queries don't work. For example, you can't parametrize the table ID.

 **Describe the solution you'd like**

When a `$variable` is the only thing in the query body, replace it with the value from the variable in Python. This is similar to how the first `%%bigquery` line works.

Cell 1

```
yourprojectid = "swast-scratch"
query = "SELECT * FROM something"
```

Cell 2

```
%%bigquery --project $yourprojectid
$query
```

 **Describe alternatives you've considered**

One could try and find any `$variable` in the query body, but this would be a breaking change.

 **Additional context**

Currently the `%%bigquery` has some logic to see if the body is a table ID or a full query. We should retain this logic, and it should happen after the replacement of the variable. Detecting if the body is only a `$variable` should be similar to the detection of if the body is only a table ID.

Feature request is from an email thread "python magics - fully parameterized query" with @minhaz and @shollyman 


Discussion:
I have some free cycles next week, is this already in progress or just assigned for now?
Thanks! :slightly_smiling_face: 

**Update:** Discussed offline, I can handle this and then move on to some `v3` issues, and can assign Lo as a reviewer for onboarding purposes.

Resolving patch:
diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py
--- a/google/cloud/bigquery/magics/magics.py
+++ b/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):


Test patch:
diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py
--- a/tests/unit/test_magics.py
+++ b/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():