From bf5964614f6ab2071ecc22c53536b09e05c3b8ea Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 23 Sep 2024 18:29:39 +0000 Subject: [PATCH 1/3] feat: add support for overriding with magic arguments --- bigquery_magics/bigquery.py | 20 ++++++++++-- tests/unit/test_bigquery.py | 62 ++++++++++++++++++++++++++++++++----- 2 files changed, 72 insertions(+), 10 deletions(-) diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py index 4cd0c5e..c7816d8 100644 --- a/bigquery_magics/bigquery.py +++ b/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 != "bigframes" or args.engine != "pandas": + raise ValueError(f"Invalid engine: {args.engine}") + + return params, args def _split_args_line(line: str) -> Tuple[str, str]: diff --git a/tests/unit/test_bigquery.py b/tests/unit/test_bigquery.py index 7079bea..5c61c18 100644 --- a/tests/unit/test_bigquery.py +++ b/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") From e38f20c18b6f3b59e377a03cfa17aabf2b5b5f5d Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Mon, 23 Sep 2024 18:46:08 +0000 Subject: [PATCH 2/3] fix arg validation logic --- bigquery_magics/bigquery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py index c7816d8..c3ce702 100644 --- a/bigquery_magics/bigquery.py +++ b/bigquery_magics/bigquery.py @@ -436,7 +436,7 @@ def _parse_magic_args(line: str) -> Tuple[List[Any], Any]: args = magic_arguments.parse_argstring(_cell_magic, rest_of_args) - if args.engine != "bigframes" or args.engine != "pandas": + if args.engine is not None and args.engine not in ("pandas", "bigframes"): raise ValueError(f"Invalid engine: {args.engine}") return params, args From 1a46ed93087ede4ea051e5c9852fd6b57089f1f8 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 24 Sep 2024 20:11:24 +0000 Subject: [PATCH 3/3] fix names --- bigquery_magics/bigquery.py | 4 ++-- tests/unit/test_bigquery.py | 14 +++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py index c3ce702..8b6d93d 100644 --- a/bigquery_magics/bigquery.py +++ b/bigquery_magics/bigquery.py @@ -393,9 +393,9 @@ def _cell_magic(line, query): return query = _validate_and_resolve_query(query, args) - use_bigframes_engine = args.engine or context.engine + engine = args.engine or context.engine - if use_bigframes_engine == "bigframes": + if engine == "bigframes": return _query_with_bigframes(query, params, args) return _query_with_pandas(query, params, args) diff --git a/tests/unit/test_bigquery.py b/tests/unit/test_bigquery.py index 5c61c18..03b61ed 100644 --- a/tests/unit/test_bigquery.py +++ b/tests/unit/test_bigquery.py @@ -1909,7 +1909,7 @@ def test_bigquery_magic_with_invalid_engine_raises_error(): @pytest.mark.usefixtures( "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context" ) -def test_big_query_magic_bigframes_set_in_context(): +def test_bigquery_magic_bigframes_set_in_context(): if bpd is None: pytest.skip("BigFrames not installed") @@ -1933,7 +1933,7 @@ def test_big_query_magic_bigframes_set_in_context(): @pytest.mark.usefixtures("ipython_interactive", "mock_credentials") -def test_big_query_magic_bigframes_set_in_args(): +def test_bigquery_magic_bigframes_set_in_args(): if bpd is None: pytest.skip("BigFrames not installed") @@ -1959,7 +1959,7 @@ def test_big_query_magic_bigframes_set_in_args(): @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(): +def test_bigquery_magic_bigframes__bigframes_is_not_installed__should_raise_error(): if bpd is not None: pytest.skip("BigFrames is installed") @@ -1974,7 +1974,7 @@ def test_big_query_magic_bigframes__bigframes_is_not_installed__should_raise_err @pytest.mark.usefixtures( "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context" ) -def test_big_query_magic_bigframes_with_params(): +def test_bigquery_magic_bigframes_with_params(): if bpd is None: pytest.skip("BigFrames not installed") @@ -2008,7 +2008,7 @@ def test_big_query_magic_bigframes_with_params(): @pytest.mark.usefixtures( "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context" ) -def test_big_query_magic_bigframes_with_max_results(): +def test_bigquery_magic_bigframes_with_max_results(): if bpd is None: pytest.skip("BigFrames not installed") @@ -2032,7 +2032,7 @@ def test_big_query_magic_bigframes_with_max_results(): @pytest.mark.usefixtures( "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context" ) -def test_big_query_magic_bigframes_with_destination_var(ipython_ns_cleanup): +def test_bigquery_magic_bigframes_with_destination_var(ipython_ns_cleanup): if bpd is None: pytest.skip("BigFrames not installed") @@ -2054,7 +2054,7 @@ def test_big_query_magic_bigframes_with_destination_var(ipython_ns_cleanup): @pytest.mark.usefixtures( "ipython_interactive", "mock_credentials", "set_bigframes_engine_in_context" ) -def test_big_query_magic_bigframes_with_dry_run__should_fail(): +def test_bigquery_magic_bigframes_with_dry_run__should_fail(): if bpd is None: pytest.skip("BigFrames not installed")