fix: detect the location of INFORMATION_SCHEMA views in read_gbq

When no location is set, bigframes.pandas.read_gbq() and
read_gbq_table() look up the table to pick the session location. That
failed for INFORMATION_SCHEMA views: tables.get can't look them up, and
4-part IDs such as project.region-us.INFORMATION_SCHEMA.SCHEMATA were
sent to the pyiceberg REST catalog. Get
the location from a dry run query on the view instead, the same way the
session loader gets metadata for these views.

BUG=570647343

Change-Id: I01b9d07e25c29f8315f54f9036393d30891e2ab9
Reviewed-on: https://bigframes-internal-review.git.corp.google.com/c/bigframes/+/4820
Kokoro-System: Kokoro <noreply+kokoro-dedicatedkokoro-dedicated Kokorogoogle.com>
Kokoro-Unit: Kokoro <noreply+kokoro-dedicatedkokoro-dedicated Kokorogoogle.com>
Kokoro-Docs: Kokoro <noreply+kokoro-dedicatedkokoro-dedicated Kokorogoogle.com>
Kokoro-E2E: Kokoro <noreply+kokoro-dedicatedkokoro-dedicated Kokorogoogle.com>
Reviewed-by: Tim Swena <swast@google.com>
Kokoro-Doctest: Kokoro <noreply+kokoro-dedicatedkokoro-dedicated Kokorogoogle.com>
Kokoro-Lint: Kokoro <noreply+kokoro-dedicatedkokoro-dedicated Kokorogoogle.com>

Author: sindhuvy

bigframes/pandas/io/api.py

@@ -687,6 +687,7 @@ def _set_default_session_location_if_possible_deferred_query(create_query):
     # Address circular imports in doctest due to bigframes/session/__init__.py
     # containing a lot of logic and samples.
     from bigframes.session._io import bigquery
+    from bigframes.session._io.bigquery import read_gbq_table as bf_read_gbq_table
 
     # Set the location as per the query if this is the first query the user is
     # running and:
@@ -715,6 +716,18 @@ def _set_default_session_location_if_possible_deferred_query(create_query):
             # aren't necessary.
             job = _dry_run(query, bqclient)
             config.options.bigquery.location = job.location
+        elif bf_read_gbq_table.is_information_schema(query):
+            # INFORMATION_SCHEMA views can't be looked up with tables.get, and
+            # 4-part IDs such as project.region-us.INFORMATION_SCHEMA.SCHEMATA
+            # look like Iceberg table IDs. Instead, dry run a query on the view
+            # and use the location of that dry run job. The session loader
+            # gets the metadata for these views the same way.
+            table = bf_read_gbq_table.get_information_schema_metadata(
+                bqclient=bqclient,
+                table_id=query,
+                default_project=default_project,
+            )
+            config.options.bigquery.location = table.location
         elif bq_data.is_irc_table(query):
             irc_table = bigframes.session.iceberg.get_table(
                 default_project, query, bqclient._credentials

tests/system/small/pandas/test_read_gbq_information_schema.py

@@ -14,6 +14,8 @@
 
 import pytest
 
+import bigframes.pandas as bpd
+
 
 @pytest.mark.parametrize("include_project", [True, False])
 @pytest.mark.parametrize(
@@ -48,3 +50,18 @@ def test_read_gbq_schemata_four_parts_can_be_peeked(unordered_session):
     )
     result = df.peek()
     assert result is not None
+
+
+@pytest.mark.parametrize("include_project", [True, False])
+def test_bpd_read_gbq_information_schema_sets_location(
+    reset_default_session_and_location, project_id: str, include_project: bool
+):
+    view_id = "region-US.INFORMATION_SCHEMA.SCHEMATA"
+    if include_project:
+        view_id = f"{project_id}.{view_id}"
+    assert not bpd.options.bigquery.location
+
+    df = bpd.read_gbq(view_id, max_results=10)
+
+    assert df.dtypes is not None
+    assert bpd.options.bigquery.location == "US"

tests/unit/pandas/io/test_api.py

@@ -127,3 +127,72 @@ def test_read_gbq_colab_calls_set_location(
     assert kwargs["pyformat_args"] == sample_pyformat_args
     assert not kwargs["dry_run"]
     assert isinstance(result, bigframes.dataframe.DataFrame)
+
+
+@pytest.mark.parametrize(
+    "table_id",
+    [
+        # 4 parts, the same shape as an Iceberg REST catalog table ID.
+        "my-project.region-us.INFORMATION_SCHEMA.SCHEMATA",
+        "my-project.my_dataset.INFORMATION_SCHEMA.TABLES",
+        # 3 parts, which tables.get would parse as project.dataset.table.
+        "region-us.INFORMATION_SCHEMA.SCHEMATA",
+        "my_dataset.INFORMATION_SCHEMA.TABLES",
+    ],
+)
+@mock.patch("bigframes.pandas.io.api._get_bqclient_and_project")
+def test_set_default_session_location_information_schema_uses_dry_run(
+    mock_get_bqclient_and_project, table_id
+):
+    bigframes.pandas.close_session()
+    bigframes.pandas.options.bigquery.location = None
+    mock_bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
+    mock_query_job = mock.create_autospec(google.cloud.bigquery.QueryJob, instance=True)
+    mock_query_job.location = "us-east4"
+    type(mock_query_job).schema = mock.PropertyMock(return_value=[])
+    mock_bqclient.query.return_value = mock_query_job
+    mock_get_bqclient_and_project.return_value = (mock_bqclient, "default-project")
+
+    try:
+        bf_io_api._set_default_session_location_if_possible(table_id)
+
+        mock_bqclient.query.assert_called_once()
+        args, kwargs = mock_bqclient.query.call_args
+        assert table_id in args[0]
+        assert kwargs["job_config"].dry_run
+        mock_bqclient.get_table.assert_not_called()
+        assert bigframes.pandas.options.bigquery.location == "us-east4"
+    finally:
+        bigframes.pandas.options.bigquery.location = None
+
+
+@pytest.mark.parametrize(
+    "table_id",
+    [
+        "my-project.my_dataset.my_table",
+        # Names that contain INFORMATION_SCHEMA but aren't INFORMATION_SCHEMA
+        # views.
+        "my-project.MY_INFORMATION_SCHEMA.TABLES",
+        "my-project.my_dataset.INFORMATION_SCHEMA",
+    ],
+)
+@mock.patch("bigframes.pandas.io.api._get_bqclient_and_project")
+def test_set_default_session_location_table_uses_get_table(
+    mock_get_bqclient_and_project, table_id
+):
+    bigframes.pandas.close_session()
+    bigframes.pandas.options.bigquery.location = None
+    mock_bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True)
+    mock_table = mock.create_autospec(google.cloud.bigquery.Table, instance=True)
+    mock_table.location = "asia-northeast1"
+    mock_bqclient.get_table.return_value = mock_table
+    mock_get_bqclient_and_project.return_value = (mock_bqclient, "default-project")
+
+    try:
+        bf_io_api._set_default_session_location_if_possible(table_id)
+
+        mock_bqclient.get_table.assert_called_once_with(table_id)
+        mock_bqclient.query.assert_not_called()
+        assert bigframes.pandas.options.bigquery.location == "asia-northeast1"
+    finally:
+        bigframes.pandas.options.bigquery.location = None