From ae4a296be4df490bc38c1e8de3589f2ace75cdd9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tim=20Swe=C3=B1a?= Date: Mon, 14 Sep 2026 16:44:25 +0000 Subject: [PATCH] fix: disambiguate google-cloud-bigquery `to_dataframe` usage from `pandas-gbq` in ua. --- .../google/cloud/bigquery/client.py | 18 +- .../tests/unit/test_client.py | 70 +++-- .../tests/unit/test_table.py | 247 +++++++++++++++++- 3 files changed, 295 insertions(+), 40 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py index 07fec2fc0fa9..f5bf674e9956 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/client.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/client.py @@ -76,6 +76,7 @@ _versions_helpers, enums, job, + version, ) from google.cloud.bigquery import exceptions as bq_exceptions from google.cloud.bigquery._helpers import ( @@ -128,9 +129,7 @@ ) pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import() -pandas = ( - _versions_helpers.PANDAS_VERSIONS.try_import() -) # mypy check fails because pandas import is outside module, there are type: ignore comments related to this +pandas = _versions_helpers.PANDAS_VERSIONS.try_import() # mypy check fails because pandas import is outside module, there are type: ignore comments related to this ResumableTimeoutType = Union[ @@ -148,7 +147,7 @@ _RESUMABLE_URL_TEMPLATE = _BASE_UPLOAD_TEMPLATE + "resumable" _GENERIC_CONTENT_TYPE = "*/*" _READ_LESS_THAN_SIZE = ( - "Size {:d} was specified but the file-like object only had " "{:d} bytes remaining." + "Size {:d} was specified but the file-like object only had {:d} bytes remaining." ) _NEED_TABLE_ARGUMENT = ( "The table argument should be a table ID string, Table, or TableReference" @@ -641,9 +640,16 @@ def _ensure_bqstorage_client( pandas_gbq = None # type: ignore if pandas_gbq is None: - user_agent = "pandas-gbq/0.0.0" + # Even if pandas-gbq isn't installed, attribute all + # to_dataframe/to_arrow usage the same as we do the recommended + # (pandas-gbq) code paths. + pandas_user_agent = "pandas-gbq/0.0.0" else: - user_agent = f"pandas-gbq/{pandas_gbq.__version__}" + pandas_user_agent = f"pandas-gbq/{pandas_gbq.__version__}" + + # Track the google-cloud-bigquery version as "legacy" because this code + # path is intended to be migrated to pandas-gbq itself. + user_agent = f"legacy-gcb/{version.__version__} {pandas_user_agent}" if client_info is None: amended_client_info = google.api_core.gapic_v1.client_info.ClientInfo( diff --git a/packages/google-cloud-bigquery/tests/unit/test_client.py b/packages/google-cloud-bigquery/tests/unit/test_client.py index 5cce574d3ff1..c48a73cc2ebd 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_client.py +++ b/packages/google-cloud-bigquery/tests/unit/test_client.py @@ -54,7 +54,7 @@ import google.cloud.bigquery.table from google.api_core import client_info from google.cloud import bigquery -from google.cloud.bigquery import ParquetOptions, exceptions +from google.cloud.bigquery import ParquetOptions, exceptions, version from google.cloud.bigquery.dataset import Dataset, DatasetReference from google.cloud.bigquery.enums import DatasetView, TimestampPrecision, UpdateMode from google.cloud.bigquery.retry import DEFAULT_TIMEOUT @@ -840,6 +840,9 @@ def test_ensure_bqstorage_client_creating_new_instance(self): self.assertIs(kwargs["client_options"], mock.sentinel.client_options) self.assertIn("test-agent", kwargs["client_info"].user_agent) self.assertIn("pandas-gbq", kwargs["client_info"].user_agent) + self.assertIn( + f"legacy-gcb/{version.__version__}", kwargs["client_info"].user_agent + ) def test_ensure_bqstorage_client_pandas_gbq_installed(self): bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") @@ -857,15 +860,17 @@ def test_ensure_bqstorage_client_pandas_gbq_installed(self): mock_pandas = mock.Mock() mock_pandas.__version__ = "0.13.0" - with mock.patch( - "google.cloud.bigquery_storage.BigQueryReadClient", mock_client - ), mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas}): + with ( + mock.patch("google.cloud.bigquery_storage.BigQueryReadClient", mock_client), + mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas}), + ): client._ensure_bqstorage_client(client_info=client_info) mock_client.assert_called_once() _, kwargs = mock_client.call_args self.assertEqual( - kwargs["client_info"].user_agent, "app-agent pandas-gbq/0.13.0" + kwargs["client_info"].user_agent, + f"app-agent legacy-gcb/{version.__version__} pandas-gbq/0.13.0", ) def test_ensure_bqstorage_client_pandas_gbq_not_installed(self): @@ -881,14 +886,18 @@ def test_ensure_bqstorage_client_pandas_gbq_not_installed(self): user_agent="app-agent" ) - with mock.patch( - "google.cloud.bigquery_storage.BigQueryReadClient", mock_client - ), mock.patch.dict(sys.modules, {"pandas_gbq": None}): + with ( + mock.patch("google.cloud.bigquery_storage.BigQueryReadClient", mock_client), + mock.patch.dict(sys.modules, {"pandas_gbq": None}), + ): client._ensure_bqstorage_client(client_info=client_info) mock_client.assert_called_once() _, kwargs = mock_client.call_args - self.assertEqual(kwargs["client_info"].user_agent, "app-agent pandas-gbq/0.0.0") + self.assertEqual( + kwargs["client_info"].user_agent, + f"app-agent legacy-gcb/{version.__version__} pandas-gbq/0.0.0", + ) def test_ensure_bqstorage_client_client_info_none(self): bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") @@ -898,14 +907,18 @@ def test_ensure_bqstorage_client_client_info_none(self): creds = _make_credentials() client = self._make_one(project=self.PROJECT, credentials=creds) - with mock.patch( - "google.cloud.bigquery_storage.BigQueryReadClient", mock_client - ), mock.patch.dict(sys.modules, {"pandas_gbq": None}): + with ( + mock.patch("google.cloud.bigquery_storage.BigQueryReadClient", mock_client), + mock.patch.dict(sys.modules, {"pandas_gbq": None}), + ): client._ensure_bqstorage_client(client_info=None) mock_client.assert_called_once() _, kwargs = mock_client.call_args - self.assertEqual(kwargs["client_info"].user_agent, "pandas-gbq/0.0.0") + self.assertEqual( + kwargs["client_info"].user_agent, + f"legacy-gcb/{version.__version__} pandas-gbq/0.0.0", + ) def test_ensure_bqstorage_client_client_info_user_agent_none(self): bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") @@ -919,14 +932,18 @@ def test_ensure_bqstorage_client_client_info_user_agent_none(self): client_info = google.api_core.gapic_v1.client_info.ClientInfo(user_agent=None) - with mock.patch( - "google.cloud.bigquery_storage.BigQueryReadClient", mock_client - ), mock.patch.dict(sys.modules, {"pandas_gbq": None}): + with ( + mock.patch("google.cloud.bigquery_storage.BigQueryReadClient", mock_client), + mock.patch.dict(sys.modules, {"pandas_gbq": None}), + ): client._ensure_bqstorage_client(client_info=client_info) mock_client.assert_called_once() _, kwargs = mock_client.call_args - self.assertEqual(kwargs["client_info"].user_agent, "pandas-gbq/0.0.0") + self.assertEqual( + kwargs["client_info"].user_agent, + f"legacy-gcb/{version.__version__} pandas-gbq/0.0.0", + ) def test_ensure_bqstorage_client_missing_dependency(self): creds = _make_credentials() @@ -9157,9 +9174,10 @@ def test_load_table_from_dataframe_w_partial_schema_extra_types(self): SchemaField("unknown_col", "BYTES"), ) job_config = job.LoadJobConfig(schema=schema) - with load_patch as load_table_from_file, pytest.raises( - ValueError - ) as exc_context: + with ( + load_patch as load_table_from_file, + pytest.raises(ValueError) as exc_context, + ): client.load_table_from_dataframe( dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION ) @@ -9429,9 +9447,13 @@ def test_load_table_from_dataframe_emits_pending_deprecation_warning(self): get_table_patch = mock.patch( "google.cloud.bigquery.client.Client.get_table", autospec=True ) - with load_patch, get_table_patch, pytest.warns( - PendingDeprecationWarning, - match="Loading DataFrames via google-cloud-bigquery is deprecated", + with ( + load_patch, + get_table_patch, + pytest.warns( + PendingDeprecationWarning, + match="Loading DataFrames via google-cloud-bigquery is deprecated", + ), ): client.load_table_from_dataframe(dataframe, self.TABLE_REF) @@ -9853,7 +9875,7 @@ def test_load_table_from_json_unicode_emoji_data_case(self): client = self._make_client() - emoji = "\U0001F3E6" + emoji = "\U0001f3e6" json_row = {"emoji": emoji} json_rows = [json_row] diff --git a/packages/google-cloud-bigquery/tests/unit/test_table.py b/packages/google-cloud-bigquery/tests/unit/test_table.py index 58612fc61f05..17719873b4f5 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_table.py +++ b/packages/google-cloud-bigquery/tests/unit/test_table.py @@ -377,9 +377,7 @@ def test_from_api_repr(self): def test___repr__(self): dataset = DatasetReference("project1", "dataset1") table1 = self._make_one(dataset, "table1") - expected = ( - "TableReference(DatasetReference('project1', 'dataset1'), " "'table1')" - ) + expected = "TableReference(DatasetReference('project1', 'dataset1'), 'table1')" self.assertEqual(repr(table1), expected) def test___str__(self): @@ -1832,9 +1830,7 @@ def test___repr__(self): dataset = DatasetReference("project1", "dataset1") table1 = self._make_one(TableReference(dataset, "table1")) expected = ( - "Table(TableReference(" - "DatasetReference('project1', 'dataset1'), " - "'table1'))" + "Table(TableReference(DatasetReference('project1', 'dataset1'), 'table1'))" ) self.assertEqual(repr(table1), expected) @@ -2894,7 +2890,8 @@ def test__should_use_bqstorage_returns_true_if_no_cached_results(self): def test__should_use_bqstorage_returns_false_if_page_token_set(self): iterator = self._make_one( - page_token="abc", first_page_response=None # not cached + page_token="abc", + first_page_response=None, # not cached ) result = iterator._should_use_bqstorage( bqstorage_client=None, create_bqstorage_client=True @@ -2903,7 +2900,8 @@ def test__should_use_bqstorage_returns_false_if_page_token_set(self): def test__should_use_bqstorage_returns_false_if_max_results_set(self): iterator = self._make_one( - max_results=10, first_page_response=None # not cached + max_results=10, + first_page_response=None, # not cached ) result = iterator._should_use_bqstorage( bqstorage_client=None, create_bqstorage_client=True @@ -3549,6 +3547,121 @@ def test_to_arrow_w_bqstorage_creates_client(self): mock_client._ensure_bqstorage_client.assert_called_once() bqstorage_client._transport.close.assert_called_once() + def test_to_arrow_create_read_session_user_agent(self): + pytest.importorskip("numpy") + pytest.importorskip("pyarrow") + pytest.importorskip("google.cloud.bigquery_storage") + import sys + + import google.auth.credentials + from google.cloud import bigquery_storage + from google.cloud.bigquery import client as client_module + from google.cloud.bigquery import schema, version + from google.cloud.bigquery import table as mut + from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( + grpc as big_query_read_grpc_transport, + ) + + mock_channel = mock.MagicMock() + mock_unary = mock.MagicMock() + mock_channel.unary_unary.return_value = mock_unary + mock_session = bigquery_storage.types.ReadSession( + name="projects/proj/locations/us/sessions/s1", + streams=[], + ) + mock_unary.with_call.return_value = (mock_session, mock.MagicMock()) + + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.__version__ = "0.13.0" + + with ( + mock.patch.object( + big_query_read_grpc_transport.BigQueryReadGrpcTransport, + "create_channel", + return_value=mock_channel, + ), + mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}), + ): + client = client_module.Client( + project="proj", + credentials=mock.Mock(spec=google.auth.credentials.Credentials), + ) + row_iterator = mut.RowIterator( + client, + None, # api_request: ignored + None, # path: ignored + [schema.SchemaField("colA", "STRING")], + table=mut.TableReference.from_string("proj.dset.tbl"), + total_rows=0, + ) + row_iterator.to_arrow(create_bqstorage_client=True) + + mock_unary.with_call.assert_called_once() + _, kwargs = mock_unary.with_call.call_args + metadata_dict = dict(kwargs["metadata"]) + self.assertIn("x-goog-api-client", metadata_dict) + user_agent = metadata_dict["x-goog-api-client"] + self.assertIn( + f"legacy-gcb/{version.__version__} pandas-gbq/0.13.0", + user_agent, + ) + + def test_to_arrow_create_read_session_user_agent_pandas_gbq_not_installed(self): + pytest.importorskip("numpy") + pytest.importorskip("pyarrow") + pytest.importorskip("google.cloud.bigquery_storage") + import sys + + import google.auth.credentials + from google.cloud import bigquery_storage + from google.cloud.bigquery import client as client_module + from google.cloud.bigquery import schema, version + from google.cloud.bigquery import table as mut + from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( + grpc as big_query_read_grpc_transport, + ) + + mock_channel = mock.MagicMock() + mock_unary = mock.MagicMock() + mock_channel.unary_unary.return_value = mock_unary + mock_session = bigquery_storage.types.ReadSession( + name="projects/proj/locations/us/sessions/s1", + streams=[], + ) + mock_unary.with_call.return_value = (mock_session, mock.MagicMock()) + + with ( + mock.patch.object( + big_query_read_grpc_transport.BigQueryReadGrpcTransport, + "create_channel", + return_value=mock_channel, + ), + mock.patch.dict(sys.modules, {"pandas_gbq": None}), + ): + client = client_module.Client( + project="proj", + credentials=mock.Mock(spec=google.auth.credentials.Credentials), + ) + row_iterator = mut.RowIterator( + client, + None, # api_request: ignored + None, # path: ignored + [schema.SchemaField("colA", "STRING")], + table=mut.TableReference.from_string("proj.dset.tbl"), + total_rows=0, + ) + row_iterator.to_arrow(create_bqstorage_client=True) + + mock_unary.with_call.assert_called_once() + _, kwargs = mock_unary.with_call.call_args + metadata_dict = dict(kwargs["metadata"]) + self.assertIn("x-goog-api-client", metadata_dict) + user_agent = metadata_dict["x-goog-api-client"] + self.assertIn( + f"legacy-gcb/{version.__version__} pandas-gbq/0.0.0", + user_agent, + ) + def test_to_arrow_ensure_bqstorage_client_wo_bqstorage(self): pytest.importorskip("numpy") pyarrow = pytest.importorskip( @@ -4904,6 +5017,121 @@ def test_to_dataframe_w_bqstorage_creates_client(self): mock_client._ensure_bqstorage_client.assert_called_once() bqstorage_client._transport.close.assert_called_once() + def test_to_dataframe_create_read_session_user_agent(self): + pytest.importorskip("numpy") + pytest.importorskip("pandas") + pytest.importorskip("google.cloud.bigquery_storage") + import sys + + import google.auth.credentials + from google.cloud import bigquery_storage + from google.cloud.bigquery import client as client_module + from google.cloud.bigquery import schema, version + from google.cloud.bigquery import table as mut + from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( + grpc as big_query_read_grpc_transport, + ) + + mock_channel = mock.MagicMock() + mock_unary = mock.MagicMock() + mock_channel.unary_unary.return_value = mock_unary + mock_session = bigquery_storage.types.ReadSession( + name="projects/proj/locations/us/sessions/s1", + streams=[], + ) + mock_unary.with_call.return_value = (mock_session, mock.MagicMock()) + + mock_pandas_gbq = mock.Mock() + mock_pandas_gbq.__version__ = "0.13.0" + + with ( + mock.patch.object( + big_query_read_grpc_transport.BigQueryReadGrpcTransport, + "create_channel", + return_value=mock_channel, + ), + mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}), + ): + client = client_module.Client( + project="proj", + credentials=mock.Mock(spec=google.auth.credentials.Credentials), + ) + row_iterator = mut.RowIterator( + client, + None, # api_request: ignored + None, # path: ignored + [schema.SchemaField("colA", "STRING")], + table=mut.TableReference.from_string("proj.dset.tbl"), + total_rows=0, + ) + row_iterator.to_dataframe(create_bqstorage_client=True) + + mock_unary.with_call.assert_called_once() + _, kwargs = mock_unary.with_call.call_args + metadata_dict = dict(kwargs["metadata"]) + self.assertIn("x-goog-api-client", metadata_dict) + user_agent = metadata_dict["x-goog-api-client"] + self.assertIn( + f"legacy-gcb/{version.__version__} pandas-gbq/0.13.0", + user_agent, + ) + + def test_to_dataframe_create_read_session_user_agent_pandas_gbq_not_installed(self): + pytest.importorskip("numpy") + pytest.importorskip("pandas") + pytest.importorskip("google.cloud.bigquery_storage") + import sys + + import google.auth.credentials + from google.cloud import bigquery_storage + from google.cloud.bigquery import client as client_module + from google.cloud.bigquery import schema, version + from google.cloud.bigquery import table as mut + from google.cloud.bigquery_storage_v1.services.big_query_read.transports import ( + grpc as big_query_read_grpc_transport, + ) + + mock_channel = mock.MagicMock() + mock_unary = mock.MagicMock() + mock_channel.unary_unary.return_value = mock_unary + mock_session = bigquery_storage.types.ReadSession( + name="projects/proj/locations/us/sessions/s1", + streams=[], + ) + mock_unary.with_call.return_value = (mock_session, mock.MagicMock()) + + with ( + mock.patch.object( + big_query_read_grpc_transport.BigQueryReadGrpcTransport, + "create_channel", + return_value=mock_channel, + ), + mock.patch.dict(sys.modules, {"pandas_gbq": None}), + ): + client = client_module.Client( + project="proj", + credentials=mock.Mock(spec=google.auth.credentials.Credentials), + ) + row_iterator = mut.RowIterator( + client, + None, # api_request: ignored + None, # path: ignored + [schema.SchemaField("colA", "STRING")], + table=mut.TableReference.from_string("proj.dset.tbl"), + total_rows=0, + ) + row_iterator.to_dataframe(create_bqstorage_client=True) + + mock_unary.with_call.assert_called_once() + _, kwargs = mock_unary.with_call.call_args + metadata_dict = dict(kwargs["metadata"]) + self.assertIn("x-goog-api-client", metadata_dict) + user_agent = metadata_dict["x-goog-api-client"] + self.assertIn( + f"legacy-gcb/{version.__version__} pandas-gbq/0.0.0", + user_agent, + ) + def test_to_dataframe_w_bqstorage_no_streams(self): pytest.importorskip("numpy") pytest.importorskip("pandas") @@ -5623,8 +5851,7 @@ def test_to_geodataframe_no_geog(self): with self.assertRaisesRegex( TypeError, re.escape( - "There must be at least one GEOGRAPHY column" - " to create a GeoDataFrame" + "There must be at least one GEOGRAPHY column to create a GeoDataFrame" ), ): row_iterator.to_geodataframe(create_bqstorage_client=False)