feat: add support for overriding context.engine by magic argument

bigquery_magics/bigquery.py

@@ -360,6 +360,15 @@ def _create_dataset_if_necessary(client, dataset_id):
         "This flag is ignored when the engine is 'bigframes'."
     ),
 )
+@magic_arguments.argument(
+    "--engine",
+    type=str,
+    default=None,
+    help=(
+        "Set the execution engine, either 'pandas' or 'bigframes'."
+        "Defaults to engine set in the query setting in console."
+    ),
+)
 def _cell_magic(line, query):
     """Underlying function for bigquery cell magic
 
@@ -384,7 +393,9 @@ def _cell_magic(line, query):
         return
     query = _validate_and_resolve_query(query, args)
 
-    if context.engine == "bigframes":
+    use_bigframes_engine = args.engine or context.engine
+
+    if use_bigframes_engine == "bigframes":
         return _query_with_bigframes(query, params, args)
 
     return _query_with_pandas(query, params, args)
@@ -423,7 +434,12 @@ def _parse_magic_args(line: str) -> Tuple[List[Any], Any]:
 
         params = _helpers.to_query_parameters(ast.literal_eval(params_option_value), {})
 
-    return params, magic_arguments.parse_argstring(_cell_magic, rest_of_args)
+    args = magic_arguments.parse_argstring(_cell_magic, rest_of_args)
+
+    if args.engine is not None and args.engine not in ("pandas", "bigframes"):
+        raise ValueError(f"Invalid engine: {args.engine}")
+
+    return params, args
 
 
 def _split_args_line(line: str) -> Tuple[str, str]:

tests/unit/test_bigquery.py

@@ -127,7 +127,7 @@ def mock_credentials(monkeypatch):
 
 
 @pytest.fixture
-def bigframes_engine(monkeypatch):
+def set_bigframes_engine_in_context(monkeypatch):
     monkeypatch.setattr(bigquery_magics.context, "engine", "bigframes")
 
 
@@ -1896,8 +1896,20 @@ def test_bigquery_magic_with_location():
         assert client_options_used.location == "us-east1"
 
 
-@pytest.mark.usefixtures("ipython_interactive", "mock_credentials", "bigframes_engine")
-def test_big_query_magic_bigframes():
+@pytest.mark.usefixtures("ipython_interactive")
+def test_bigquery_magic_with_invalid_engine_raises_error():
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    engine = "whatever"
+
+    with pytest.raises(ValueError, match=f"Invalid engine: {engine}"):
+        ip.run_cell_magic("bigquery", f"--engine {engine}", "SELECT 17 AS num")
+
+
+@pytest.mark.usefixtures(
+    "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context"
+)
+def test_big_query_magic_bigframes_set_in_context():
     if bpd is None:
         pytest.skip("BigFrames not installed")
 
@@ -1920,7 +1932,33 @@ def test_big_query_magic_bigframes():
         assert bpd.options.bigquery.project == bigquery_magics.context.project
 
 
-@pytest.mark.usefixtures("ipython_interactive", "mock_credentials", "bigframes_engine")
+@pytest.mark.usefixtures("ipython_interactive", "mock_credentials")
+def test_big_query_magic_bigframes_set_in_args():
+    if bpd is None:
+        pytest.skip("BigFrames not installed")
+
+    ip = IPython.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    sql = "SELECT 0 AS something"
+    expected_configuration = {
+        "query": {"queryParameters": [], "useLegacySql": False},
+        "dryRun": False,
+    }
+    bf_patch = mock.patch("bigframes.pandas.read_gbq_query", autospec=True)
+
+    with bf_patch as bf_mock:
+        ip.run_cell_magic("bigquery", "--engine bigframes", sql)
+
+        bf_mock.assert_called_once_with(
+            sql, max_results=None, configuration=expected_configuration
+        )
+        assert bpd.options.bigquery.credentials is bigquery_magics.context.credentials
+        assert bpd.options.bigquery.project == bigquery_magics.context.project
+
+
+@pytest.mark.usefixtures(
+    "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context"
+)
 def test_big_query_magic_bigframes__bigframes_is_not_installed__should_raise_error():
     if bpd is not None:
         pytest.skip("BigFrames is installed")
@@ -1933,7 +1971,9 @@ def test_big_query_magic_bigframes__bigframes_is_not_installed__should_raise_err
         ip.run_cell_magic("bigquery", "", sql)
 
 
-@pytest.mark.usefixtures("ipython_interactive", "mock_credentials", "bigframes_engine")
+@pytest.mark.usefixtures(
+    "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context"
+)
 def test_big_query_magic_bigframes_with_params():
     if bpd is None:
         pytest.skip("BigFrames not installed")
@@ -1965,7 +2005,9 @@ def test_big_query_magic_bigframes_with_params():
         )
 
 
-@pytest.mark.usefixtures("ipython_interactive", "mock_credentials", "bigframes_engine")
+@pytest.mark.usefixtures(
+    "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context"
+)
 def test_big_query_magic_bigframes_with_max_results():
     if bpd is None:
         pytest.skip("BigFrames not installed")
@@ -1987,7 +2029,9 @@ def test_big_query_magic_bigframes_with_max_results():
         )
 
 
-@pytest.mark.usefixtures("ipython_interactive", "mock_credentials", "bigframes_engine")
+@pytest.mark.usefixtures(
+    "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context"
+)
 def test_big_query_magic_bigframes_with_destination_var(ipython_ns_cleanup):
     if bpd is None:
         pytest.skip("BigFrames not installed")
@@ -2007,7 +2051,9 @@ def test_big_query_magic_bigframes_with_destination_var(ipython_ns_cleanup):
         assert df is bf_mock.return_value
 
 
-@pytest.mark.usefixtures("ipython_interactive", "mock_credentials", "bigframes_engine")
+@pytest.mark.usefixtures(
+    "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context"
+)
 def test_big_query_magic_bigframes_with_dry_run__should_fail():
     if bpd is None:
         pytest.skip("BigFrames not installed")