Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions google/cloud/bigquery/magics/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@
A dataset and table to store the query results. If table does not exists,
it will be created. If table already exists, its data will be overwritten.
Variable should be in a format <dataset_id>.<table_id>.
* ``--no_query_cache`` (Optional[line argument]):
Do not use cached query results.
* ``--project <project>`` (Optional[line argument]):
Project to use for running the query. Defaults to the context
:attr:`~google.cloud.bigquery.magics.Context.project`.
Expand Down Expand Up @@ -442,6 +444,12 @@ def _create_dataset_if_necessary(client, dataset_id):
"this option's value in the context bqstorage_client_options."
),
)
@magic_arguments.argument(
"--no_query_cache",
action="store_true",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not understanding what this "store_true" is about, but it seems consistent with the other arguments.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's behavior IPython/Jupyter inherits from argparse https://docs.python.org/3/library/argparse.html#action

The gist is that if this argument is present it'll set args.no_query_cache = True rather than requiring a boolean parameter.

default=False,
help=("Do not use cached query results."),
)
@magic_arguments.argument(
"--use_bqstorage_api",
action="store_true",
Expand Down Expand Up @@ -642,6 +650,10 @@ def _cell_magic(line, query):
job_config.use_legacy_sql = args.use_legacy_sql
job_config.dry_run = args.dry_run

# Don't override context job config unless --no_query_cache is explicitly set.
if args.no_query_cache:
job_config.use_query_cache = False

if args.destination_table:
split = args.destination_table.split(".")
if len(split) != 2:
Expand Down
58 changes: 58 additions & 0 deletions tests/unit/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -1217,6 +1217,64 @@ def test_bigquery_magic_w_maximum_bytes_billed_w_context_setter():
assert sent_config["maximumBytesBilled"] == "10203"


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_with_no_query_cache(monkeypatch):
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
conn = make_connection()
monkeypatch.setattr(magics.context, "_connection", conn)
monkeypatch.setattr(magics.context, "project", "project-from-context")

# --no_query_cache option should override context.
monkeypatch.setattr(
magics.context.default_query_job_config, "use_query_cache", True
)

ip.run_cell_magic("bigquery", "--no_query_cache", QUERY_STRING)

conn.api_request.assert_called_with(
method="POST",
path="/projects/project-from-context/jobs",
data=mock.ANY,
timeout=DEFAULT_TIMEOUT,
)
jobs_insert_call = [
call
for call in conn.api_request.call_args_list
if call[1]["path"] == "/projects/project-from-context/jobs"
][0]
assert not jobs_insert_call[1]["data"]["configuration"]["query"]["useQueryCache"]


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_context_with_no_query_cache_from_context(monkeypatch):
ip = IPython.get_ipython()
ip.extension_manager.load_extension("google.cloud.bigquery")
conn = make_connection()
monkeypatch.setattr(magics.context, "_connection", conn)
monkeypatch.setattr(magics.context, "project", "project-from-context")
monkeypatch.setattr(
magics.context.default_query_job_config, "use_query_cache", False
)

ip.run_cell_magic("bigquery", "", QUERY_STRING)

conn.api_request.assert_called_with(
method="POST",
path="/projects/project-from-context/jobs",
data=mock.ANY,
timeout=DEFAULT_TIMEOUT,
)
jobs_insert_call = [
call
for call in conn.api_request.call_args_list
if call[1]["path"] == "/projects/project-from-context/jobs"
][0]
assert not jobs_insert_call[1]["data"]["configuration"]["query"]["useQueryCache"]


@pytest.mark.usefixtures("ipython_interactive")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_bigquery_magic_w_progress_bar_type_w_context_setter(monkeypatch):
Expand Down