Skip to content

Commit

Permalink
feat: allow cell magic body to be a $variable
Browse files Browse the repository at this point in the history
  • Loading branch information
plamut committed Nov 8, 2021
1 parent 4616cd5 commit c16fde3
Show file tree
Hide file tree
Showing 2 changed files with 131 additions and 1 deletion.
19 changes: 19 additions & 0 deletions google/cloud/bigquery/magics/magics.py
Expand Up @@ -596,6 +596,25 @@ 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 query_var_name.isidentifier():
ip = IPython.get_ipython()
try:
query = ip.user_ns[query_var_name]
except KeyError:
raise NameError(
f"Unknown query, variable {query_var_name} does not exist."
) from None
else:
if not isinstance(query, (str, bytes)):
raise TypeError(
f"Query variable {query_var_name} must be string "
"or 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):
Expand Down
113 changes: 112 additions & 1 deletion tests/unit/test_magics.py
Expand Up @@ -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(
Expand Down Expand Up @@ -1710,6 +1710,117 @@ 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"
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"

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
)

cell_body = "$"

with io.capture_output() as captured_io:
ip.run_cell_magic("bigquery", "", cell_body)

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_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"

with pytest.raises(
TypeError, match=r".*must be string or 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_with_invalid_multiple_option_values():
Expand Down

0 comments on commit c16fde3

Please sign in to comment.