From 7dd30af41b8a595b96176c964ba14aa41645ef0d Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Wed, 16 Feb 2022 17:07:47 -0600
Subject: [PATCH] feat: add `--no_query_cache` option to `%%bigquery` magics to
 disable query cache (#1141)

---
 google/cloud/bigquery/magics/magics.py | 12 ++++++
 tests/unit/test_magics.py              | 58 ++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/google/cloud/bigquery/magics/magics.py b/google/cloud/bigquery/magics/magics.py
index 5af0a3b51..7b4d584fb 100644
--- a/google/cloud/bigquery/magics/magics.py
+++ b/google/cloud/bigquery/magics/magics.py
@@ -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`.
@@ -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",
@@ -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:
diff --git a/tests/unit/test_magics.py b/tests/unit/test_magics.py
index e18d04d64..2801768f8 100644
--- a/tests/unit/test_magics.py
+++ b/tests/unit/test_magics.py
@@ -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):