feat: add --pyformat option to use Python string formatting (#111)

* feat: add --pyformat option to use Python string formatting

* use mock_credentials in magics tests

* avoid typing issue in tests

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* avoid polluting global options with context tests

* fix coverage

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: tswast

bigquery_magics/bigquery.py

@@ -77,6 +77,14 @@
           serializable. The variable reference is indicated by a ``$`` before
           the variable name (ex. ``$my_dict_var``). See ``In[6]`` and ``In[7]``
           in the Examples section below.
+    * ``--engine <engine>`` (Optional[line argument]):
+          Set the execution engine, either 'pandas' (default) or 'bigframes'
+          (experimental).
+    * ``--pyformat`` (Optional[line argument]):
+          Warning! Do not use with user-provided values.
+          This doesn't escape values. Use --params instead for proper SQL escaping.
+          This enables Python string formatting in the query text.
+          Useful for values not supported by SQL query params such as table IDs.
 
     * ``<query>`` (required, cell argument):
         SQL query to run. If the query does not contain any whitespace (aside
@@ -122,7 +130,7 @@
 import bigquery_magics._versions_helpers
 import bigquery_magics.config
 import bigquery_magics.graph_server as graph_server
-import bigquery_magics.line_arg_parser.exceptions
+import bigquery_magics.pyformat
 import bigquery_magics.version
 
 try:
@@ -401,6 +409,17 @@ def _create_dataset_if_necessary(client, dataset_id):
     default=False,
     help=("Visualizes the query results as a graph"),
 )
+@magic_arguments.argument(
+    "--pyformat",
+    action="store_true",
+    default=False,
+    help=(
+        "Warning! Do not use with user-provided values. "
+        "This doesn't escape values. Use --params instead for proper SQL escaping. "
+        "This enables Python string formatting in the query text. "
+        "Useful for values not supported by SQL query params such as table IDs. "
+    ),
+)
 def _cell_magic(line, query):
     """Underlying function for bigquery cell magic
 
@@ -515,7 +534,6 @@ def _query_with_bigframes(query: str, params: List[Any], args: Any):
 
 def _query_with_pandas(query: str, params: List[Any], args: Any):
     bq_client, bqstorage_client = _create_clients(args)
-
     try:
         return _make_bq_query(
             query,
@@ -779,26 +797,30 @@ def _make_bq_query(
 
 def _validate_and_resolve_query(query: str, args: Any) -> str:
     # Check if query is given as a reference to a variable.
-    if not query.startswith("$"):
-        return query
+    if query.startswith("$"):
+        query_var_name = query[1:]
 
-    query_var_name = query[1:]
+        if not query_var_name:
+            missing_msg = 'Missing query variable name, empty "$" is not allowed.'
+            raise NameError(missing_msg)
 
-    if not query_var_name:
-        missing_msg = 'Missing query variable name, empty "$" is not allowed.'
-        raise NameError(missing_msg)
+        if query_var_name.isidentifier():
+            ip = get_ipython()
+            query = ip.user_ns.get(query_var_name, ip)  # ip serves as a sentinel
 
-    if query_var_name.isidentifier():
+            if query is ip:
+                raise NameError(
+                    f"Unknown query, variable {query_var_name} does not exist."
+                )
+            elif not isinstance(query, (str, bytes)):
+                raise TypeError(
+                    f"Query variable {query_var_name} must be a string "
+                    "or a bytes-like value."
+                )
+
+    if args.pyformat:
         ip = get_ipython()
-        query = ip.user_ns.get(query_var_name, ip)  # ip serves as a sentinel
-
-        if query is ip:
-            raise NameError(f"Unknown query, variable {query_var_name} does not exist.")
-        elif not isinstance(query, (str, bytes)):
-            raise TypeError(
-                f"Query variable {query_var_name} must be a string "
-                "or a bytes-like value."
-            )
+        query = bigquery_magics.pyformat.pyformat(query, ip.user_ns)
     return query
 
 

bigquery_magics/pyformat.py

@@ -0,0 +1,77 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helpers for the --pyformat feature."""
+
+from __future__ import annotations
+
+import numbers
+import string
+from typing import Any
+
+
+def _field_to_template_value(name: str, value: Any) -> Any:
+    """Convert value to something embeddable in a SQL string.
+
+    Does **not** escape strings.
+    """
+    _validate_type(name, value)
+    # TODO(tswast): convert DataFrame objects to gbq tables or literals subquery.
+    return value
+
+
+def _validate_type(name: str, value: Any):
+    """Raises TypeError if value is unsupported."""
+    if not isinstance(value, (str, numbers.Real)):
+        raise TypeError(
+            f"{name} has unsupported type: {type(value)}. "
+            "Only str, int, float are supported."
+        )
+
+
+def _parse_fields(sql_template: str) -> list[str]:
+    return [
+        field_name
+        for _, field_name, _, _ in string.Formatter().parse(sql_template)
+        if field_name is not None
+    ]
+
+
+def pyformat(sql_template: str, user_ns: dict) -> str:
+    """Unsafe Python-style string formatting of SQL string.
+
+    Only some data types supported.
+
+    Warning: strings are **not** escaped. This allows them to be used in
+    contexts such as table identifiers, where normal query parameters are not
+    supported.
+
+    Args:
+        sql_template (str):
+            SQL string with 0+ {var_name}-style format options.
+        user_ns (dict):
+            IPython variable namespace to use for formatting.
+
+    Raises:
+        TypeError: if a referenced variable is not of a supported type.
+        KeyError: if a referenced variable is not found.
+    """
+    fields = _parse_fields(sql_template)
+
+    format_kwargs = {}
+    for name in fields:
+        value = user_ns[name]
+        format_kwargs[name] = _field_to_template_value(name, value)
+
+    return sql_template.format(**format_kwargs)

tests/system/test_bigquery.py

@@ -16,38 +16,17 @@
 
 import re
 
+from IPython.testing import globalipapp
+from IPython.utils import io
+import pandas
 import psutil
-import pytest
 
-IPython = pytest.importorskip("IPython")
-io = pytest.importorskip("IPython.utils.io")
-pandas = pytest.importorskip("pandas")
-tools = pytest.importorskip("IPython.testing.tools")
-interactiveshell = pytest.importorskip("IPython.terminal.interactiveshell")
 
-
-@pytest.fixture(scope="session")
-def ipython():
-    config = tools.default_config()
-    config.TerminalInteractiveShell.simple_prompt = True
-    shell = interactiveshell.TerminalInteractiveShell.instance(config=config)
-    return shell
-
-
-@pytest.fixture()
-def ipython_interactive(ipython):
-    """Activate IPython's builtin hooks
-
-    for the duration of the test scope.
-    """
-    with ipython.builtin_trap:
-        yield ipython
-
-
-def test_bigquery_magic(ipython_interactive):
-    ip = IPython.get_ipython()
+def test_bigquery_magic():
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
     current_process = psutil.Process()
-    conn_count_start = len(current_process.connections())
+    conn_count_start = len(current_process.net_connections())
 
     ip.extension_manager.load_extension("bigquery_magics")
     sql = """
@@ -64,7 +43,7 @@ def test_bigquery_magic(ipython_interactive):
     with io.capture_output() as captured:
         result = ip.run_cell_magic("bigquery", "--use_rest_api", sql)
 
-    conn_count_end = len(current_process.connections())
+    conn_count_end = len(current_process.net_connections())
 
     lines = re.split("\n|\r", captured.stdout)
     # Removes blanks & terminal code (result of display clearing)

tests/unit/bigquery/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unit/bigquery/conftest.py

@@ -0,0 +1,77 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+import google.auth.credentials
+import pytest
+import test_utils.imports  # google-cloud-testutils
+
+import bigquery_magics
+
+
+@pytest.fixture()
+def ipython_ns_cleanup():
+    """A helper to clean up user namespace after the test
+
+    for the duration of the test scope.
+    """
+    names_to_clean = []  # pairs (IPython_instance, name_to_clean)
+
+    yield names_to_clean
+
+    for ip, name in names_to_clean:
+        if name in ip.user_ns:
+            del ip.user_ns[name]
+
+
+@pytest.fixture(scope="session")
+def missing_bq_storage():
+    """Provide a patcher that can make the bigquery storage import to fail."""
+
+    def fail_if(name, globals, locals, fromlist, level):
+        # NOTE: *very* simplified, assuming a straightforward absolute import
+        return "bigquery_storage" in name or (
+            fromlist is not None and "bigquery_storage" in fromlist
+        )
+
+    return test_utils.imports.maybe_fail_import(predicate=fail_if)
+
+
+@pytest.fixture(scope="session")
+def missing_grpcio_lib():
+    """Provide a patcher that can make the gapic library import to fail."""
+
+    def fail_if(name, globals, locals, fromlist, level):
+        # NOTE: *very* simplified, assuming a straightforward absolute import
+        return "gapic_v1" in name or (fromlist is not None and "gapic_v1" in fromlist)
+
+    return test_utils.imports.maybe_fail_import(predicate=fail_if)
+
+
+@pytest.fixture
+def mock_credentials(monkeypatch):
+    credentials = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+
+    # Set up the context with monkeypatch so that it's reset for subsequent
+    # tests.
+    monkeypatch.setattr(bigquery_magics.context, "_project", "test-project")
+    monkeypatch.setattr(bigquery_magics.context, "_credentials", credentials)
+
+
+@pytest.fixture
+def set_bigframes_engine_in_context(monkeypatch):
+    monkeypatch.setattr(bigquery_magics.context, "engine", "bigframes")