Skip to content

Commit

Permalink
feat: add --no_query_cache option to %%bigquery magics to disable…
Browse files Browse the repository at this point in the history
… query cache (#1141)
  • Loading branch information
tswast authored Feb 16, 2022
1 parent f5daa9b commit 7dd30af
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 0 deletions.
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",
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

0 comments on commit 7dd30af

Please sign in to comment.