-
Notifications
You must be signed in to change notification settings - Fork 309
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: allow cell magic body to be a $variable #1053
Merged
+161
−1
Merged
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
c16fde3
feat: allow cell magic body to be a $variable
plamut e5610d6
Fix missing indefinitive article in error msg
plamut e9fd366
Adjust test assertion to error message change
plamut ebd3228
Merge branch 'main' into iss-1040
plamut 6c170e4
Merge branch 'main' into iss-1040
plamut e26e0a8
Refactor logic for extracting query variable
plamut 88e1638
Explicitly warn about missing query variable name
plamut b6d43c1
Thest the query "variable" is not identifier case
plamut File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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,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 = "$" | ||
plamut marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
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 a string or a bytes-like.*" | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Might be good to have another comment here to match it to the other |
||
), 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(): | ||
|
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe comment to indicate we're testing the first kind of error, as I got the regex string match from context of the other file.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The name of the test explains which case it covers ("non-existing query variable"), do you think that's not sufficient by itself?
I always look at the name of the test first, although it might be just how my brain is wired...
Update: Added a comment next to the assignment to
cell_body
to explain which use case it targets.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM - I really like how you added the "non-existing query variable" condition.