feat: add --no_query_cache option to %%bigquery magics to disable query cache (#1141)

Author: tswast

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:

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):