chore: prevent location warning when _read_gbq_colab can determine the location (#1802)

* chore: prevent location warning when _read_gbq_colab can determine the location

I've made some updates to prevent a location warning when the system can determine the location for `bigframes.pandas.io.api._read_gbq_colab`.

I've updated `bigframes.pandas.io.api._read_gbq_colab` so it correctly calls `bigframes.session.Session._read_gbq_colab` and adjusted its arguments.

The `_read_gbq_colab` function in the pandas API layer now has a simpler signature, accepting `query_or_table`, `pyformat_args`, and `dry_run`. It will continue to call `_set_default_session_location_if_possible` to prevent location warnings.

I've also updated the unit tests to reflect these changes, making sure that the correct session-level function is called and that arguments are passed through as expected. I've also moved the tests to `tests/unit/pandas/io/test_api.py` and converted them to pytest style to follow our repository conventions.

* Update tests/unit/pandas/io/test_api.py

* remove unused import

* format the query first

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: tswast

bigframes/pandas/io/api.py

@@ -216,6 +216,67 @@ def read_gbq(
 read_gbq.__doc__ = inspect.getdoc(bigframes.session.Session.read_gbq)
 
 
+@overload
+def _read_gbq_colab(  # type: ignore[overload-overlap]
+    query_or_table: str,
+    *,
+    pyformat_args: Optional[Dict[str, Any]] = ...,
+    dry_run: Literal[False] = ...,
+) -> bigframes.dataframe.DataFrame:
+    ...
+
+
+@overload
+def _read_gbq_colab(
+    query_or_table: str,
+    *,
+    pyformat_args: Optional[Dict[str, Any]] = ...,
+    dry_run: Literal[True] = ...,
+) -> pandas.Series:
+    ...
+
+
+def _read_gbq_colab(
+    query_or_table: str,
+    *,
+    pyformat_args: Optional[Dict[str, Any]] = None,
+    dry_run: bool = False,
+) -> bigframes.dataframe.DataFrame | pandas.Series:
+    """A Colab-specific version of read_gbq.
+
+    Calls `_set_default_session_location_if_possible` and then delegates
+    to `bigframes.session.Session._read_gbq_colab`.
+
+    Args:
+        query_or_table (str):
+            SQL query or table ID (table ID not yet supported).
+        pyformat_args (Optional[Dict[str, Any]]):
+            Parameters to format into the query string.
+        dry_run (bool):
+            If True, estimates the query results size without returning data.
+            The return will be a pandas Series with query metadata.
+
+    Returns:
+        Union[bigframes.dataframe.DataFrame, pandas.Series]:
+            A BigQuery DataFrame if `dry_run` is False, otherwise a pandas Series.
+    """
+    if pyformat_args is None:
+        pyformat_args = {}
+
+    query = bigframes.core.pyformat.pyformat(
+        query_or_table,
+        pyformat_args=pyformat_args,
+    )
+    _set_default_session_location_if_possible(query)
+
+    return global_session.with_default_session(
+        bigframes.session.Session._read_gbq_colab,
+        query_or_table,
+        pyformat_args=pyformat_args,
+        dry_run=dry_run,
+    )
+
+
 def read_gbq_model(model_name: str):
     return global_session.with_default_session(
         bigframes.session.Session.read_gbq_model,

tests/unit/pandas/io/test_api.py

@@ -0,0 +1,48 @@
+# Copyright 2024 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.
+
+from unittest import mock
+
+import bigframes.dataframe
+import bigframes.pandas.io.api as bf_io_api
+import bigframes.session
+
+
+@mock.patch("bigframes.pandas.io.api._set_default_session_location_if_possible")
+@mock.patch("bigframes.core.global_session.with_default_session")
+def test_read_gbq_colab_calls_set_location(
+    mock_with_default_session, mock_set_location
+):
+    # Configure the mock for with_default_session to return a DataFrame mock
+    mock_df = mock.create_autospec(bigframes.dataframe.DataFrame)
+    mock_with_default_session.return_value = mock_df
+
+    query_or_table = "SELECT {param1} AS param1"
+    sample_pyformat_args = {"param1": "value1"}
+    result = bf_io_api._read_gbq_colab(
+        query_or_table, pyformat_args=sample_pyformat_args, dry_run=False
+    )
+
+    # Make sure that we format the SQL first to prevent syntax errors.
+    formatted_query = "SELECT 'value1' AS param1"
+    mock_set_location.assert_called_once_with(formatted_query)
+    mock_with_default_session.assert_called_once()
+
+    # Check the actual arguments passed to with_default_session
+    args, kwargs = mock_with_default_session.call_args
+    assert args[0] == bigframes.session.Session._read_gbq_colab
+    assert args[1] == query_or_table
+    assert kwargs["pyformat_args"] == sample_pyformat_args
+    assert not kwargs["dry_run"]
+    assert isinstance(result, bigframes.dataframe.DataFrame)