Skip to content

feat: add user-agent delegation telemetry for pandas-gbq #17704

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
shuoweil wants to merge 15 commits into from
+247 −16
Closed

feat: add user-agent delegation telemetry for pandas-gbq #17704

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
67a61c6
fix(bigquery): harden test fixtures and prevent socket leaks
shuoweil Jul 30, 2026
f640ded
feat(bigquery): add pandas-gbq capability helper
shuoweil Jul 30, 2026
d8df911
Update packages/google-cloud-bigquery/google/cloud/bigquery/dbapi/con…
shuoweil Jul 30, 2026
1600294
Update packages/google-cloud-bigquery/google/cloud/bigquery/_versions…
shuoweil Jul 30, 2026
07539ae
Update packages/google-cloud-bigquery/google/cloud/bigquery/_versions…
shuoweil Jul 30, 2026
16ad860
style: fix line length in _versions_helpers.py
shuoweil Jul 30, 2026
07a0f85
feat(bigquery): inject user-agent telemetry for pandas-gbq delegation
shuoweil Jul 30, 2026
623276b
fix(bigquery): remove unexercised try-except in Connection.close
shuoweil Jul 30, 2026
9023a7b
chore: sync PR 1 fixes into PR 3
shuoweil Jul 30, 2026
9222108
Merge branch 'main' into shuowei-gbq-delegation-telemetry-mvp
shuoweil Aug 3, 2026
ee9ceb9
chore: sync with main and PR 2 version helper
shuoweil Aug 3, 2026
5d80f9a
feat(bigquery): harden telemetry injection and add comprehensive tests
shuoweil Aug 3, 2026
5b7db9e
Merge branch 'main' into shuowei-gbq-delegation-telemetry-mvp
shuoweil Aug 5, 2026
0ca751c
style(bigquery): chain context managers and format AAA in test_table
shuoweil Aug 5, 2026
482bb54
style(bigquery): apply black formatting to test_table
shuoweil Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 29 additions & 16 deletions packages/google-cloud-bigquery/google/cloud/bigquery/table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,7 @@
import functools
import operator
import typing
from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union, Sequence

from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple, Union
import warnings

try:
Expand Down Expand Up @@ -57,29 +56,34 @@

import google.api_core.exceptions
from google.api_core.page_iterator import HTTPIterator

import google.cloud._helpers # type: ignore
from google.cloud.bigquery import _helpers
from google.cloud.bigquery import _pandas_helpers
from google.cloud.bigquery import _versions_helpers

from google.cloud.bigquery import (
_helpers,
_pandas_helpers,
_string_references,
_versions_helpers,
)
from google.cloud.bigquery import exceptions as bq_exceptions
from google.cloud.bigquery import external_config
from google.cloud.bigquery import schema as _schema
from google.cloud.bigquery._tqdm_helpers import get_progress_bar
from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
from google.cloud.bigquery.enums import DefaultPandasDTypes
from google.cloud.bigquery.external_config import ExternalConfig
from google.cloud.bigquery import schema as _schema
from google.cloud.bigquery.schema import _build_schema_resource
from google.cloud.bigquery.schema import _parse_schema_resource
from google.cloud.bigquery.schema import _to_schema_fields
from google.cloud.bigquery import external_config
from google.cloud.bigquery import _string_references
from google.cloud.bigquery.schema import (
_build_schema_resource,
_parse_schema_resource,
_to_schema_fields,
)

if typing.TYPE_CHECKING: # pragma: NO COVER
# Unconditionally import optional dependencies again to tell pytype that
# they are not None, avoiding false "no attribute" errors.
import geopandas # type: ignore
import pandas
import pyarrow
import geopandas # type: ignore

from google.cloud import bigquery_storage # type: ignore
from google.cloud.bigquery.dataset import DatasetReference

Expand Down Expand Up @@ -797,7 +801,7 @@ def time_partitioning(self, value):
api_repr = value.to_api_repr()
elif value is not None:
raise ValueError(
"value must be google.cloud.bigquery.table.TimePartitioning " "or None"
"value must be google.cloud.bigquery.table.TimePartitioning or None"
)
self._properties[self._PROPERTY_TO_API_FIELD["time_partitioning"]] = api_repr

Expand Down Expand Up @@ -2801,6 +2805,16 @@ def to_dataframe(
create_bqstorage_client = False
bqstorage_client = None

if _versions_helpers.PANDAS_GBQ_VERSIONS.is_delegation_supported:
client_info = getattr(
getattr(self.client, "_connection", None), "_client_info", None
)
if client_info:
ua = client_info.user_agent or ""
if "pandas-gbq" not in ua:
version = _versions_helpers.PANDAS_GBQ_VERSIONS.installed_version
client_info.user_agent = f"{ua} pandas-gbq/{version}".strip()

@sycai sycai Aug 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We haven't started the pandas-gbq migration, right? My concern is that, if we start setting the user agents to be pandas-gbq without actually using that package, then our metrics will be inaccurate.

Should we hold off on this change until we start re-directing traffics to pandas-gbq?

@shuoweil shuoweil Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, @sycai! You're completely right.

While this block is currently guarded behind _versions_helpers.PANDAS_GBQ_VERSIONS.is_delegation_supported (which checks pandas_gbq._internal_delegation_api_version >= 1.0.0 and currently evaluates to False), RowIterator.to_dataframe() in google-cloud-bigquery has not yet switched over to delegating DataFrame conversion to pandas-gbq.

I will close this RP for now. We will reimplement this once we have other features ready.


record_batch = self.to_arrow(
progress_bar_type=progress_bar_type,
bqstorage_client=bqstorage_client,
Expand Down Expand Up @@ -2990,8 +3004,7 @@ def to_geodataframe(
)
if not geography_columns:
raise TypeError(
"There must be at least one GEOGRAPHY column"
" to create a GeoDataFrame"
"There must be at least one GEOGRAPHY column to create a GeoDataFrame"
)

if geography_column:
Expand Down
218 changes: 218 additions & 0 deletions packages/google-cloud-bigquery/tests/unit/test_table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import datetime
import logging
import re
import sys
import time
import types
import unittest
Expand Down Expand Up @@ -5723,6 +5724,223 @@ def test_rowiterator_to_geodataframe_delegation(self, to_dataframe):

self.assertEqual([v.__class__.__name__ for v in df.g], ["Point"])

def test_to_dataframe_delegated_updates_user_agent(self):
pytest.importorskip("db_dtypes")
pandas = pytest.importorskip("pandas")
mock_pandas_gbq = mock.Mock()
mock_pandas_gbq.__version__ = "1.0.0"

mock_client_info = mock.Mock()
mock_client_info.user_agent = "gl-python/3.10.0"

mock_client = _mock_client()
mock_client._connection = mock.Mock(_client_info=mock_client_info)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=True,
),
mock.patch(
"google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW",
False,
),
mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}),
):
row_iterator = self._make_one_from_data((("name", "STRING"),), (("foo",),))
row_iterator.client = mock_client

df = row_iterator.to_dataframe(progress_bar_type="tqdm", timeout=5.0)

self.assertIsInstance(df, pandas.DataFrame)
self.assertEqual(
mock_client_info.user_agent,
"gl-python/3.10.0 pandas-gbq/1.0.0",
)

def test_to_dataframe_delegated_does_not_duplicate_user_agent(self):
pytest.importorskip("db_dtypes")
pandas = pytest.importorskip("pandas")
mock_pandas_gbq = mock.Mock()
mock_pandas_gbq.__version__ = "1.0.0"

mock_client_info = mock.Mock()
mock_client_info.user_agent = "gl-python/3.10.0 pandas-gbq/1.0.0"

mock_client = _mock_client()
mock_client._connection = mock.Mock(_client_info=mock_client_info)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=True,
),
mock.patch(
"google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW",
False,
),
mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}),
):
row_iterator = self._make_one_from_data((("name", "STRING"),), (("foo",),))
row_iterator.client = mock_client

df = row_iterator.to_dataframe(progress_bar_type="tqdm", timeout=5.0)

self.assertIsInstance(df, pandas.DataFrame)
self.assertEqual(
mock_client_info.user_agent,
"gl-python/3.10.0 pandas-gbq/1.0.0",
)

def test_to_dataframe_delegated_when_client_info_is_none(self):
pytest.importorskip("db_dtypes")
pandas = pytest.importorskip("pandas")
mock_pandas_gbq = mock.Mock()
mock_pandas_gbq.__version__ = "1.0.0"

mock_client = _mock_client()
mock_client._connection = mock.Mock(_client_info=None)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=True,
),
mock.patch(
"google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW",
False,
),
mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}),
):
row_iterator = self._make_one_from_data((("name", "STRING"),), (("foo",),))
row_iterator.client = mock_client

df = row_iterator.to_dataframe(progress_bar_type="tqdm", timeout=5.0)

self.assertIsInstance(df, pandas.DataFrame)

def test_to_dataframe_delegated_when_user_agent_is_none(self):
pytest.importorskip("db_dtypes")
pandas = pytest.importorskip("pandas")
mock_pandas_gbq = mock.Mock()
mock_pandas_gbq.__version__ = "1.0.0"

mock_client_info = mock.Mock()
mock_client_info.user_agent = None

mock_client = _mock_client()
mock_client._connection = mock.Mock(_client_info=mock_client_info)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=True,
),
mock.patch(
"google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW",
False,
),
mock.patch.dict(sys.modules, {"pandas_gbq": mock_pandas_gbq}),
):
row_iterator = self._make_one_from_data((("name", "STRING"),), (("foo",),))
row_iterator.client = mock_client

df = row_iterator.to_dataframe(progress_bar_type="tqdm", timeout=5.0)

self.assertIsInstance(df, pandas.DataFrame)
self.assertEqual(
mock_client_info.user_agent,
"pandas-gbq/1.0.0",
)

def test_to_dataframe_delegated_false_does_not_update_user_agent(self):
pytest.importorskip("db_dtypes")
pandas = pytest.importorskip("pandas")

mock_client_info = mock.Mock()
mock_client_info.user_agent = "gl-python/3.10.0"

mock_client = _mock_client()
mock_client._connection = mock.Mock(_client_info=mock_client_info)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=False,
),
mock.patch(
"google.cloud.bigquery._versions_helpers.SUPPORTS_RANGE_PYARROW",
False,
),
):
row_iterator = self._make_one_from_data((("name", "STRING"),), (("foo",),))
row_iterator.client = mock_client

df = row_iterator.to_dataframe(progress_bar_type="tqdm", timeout=5.0)

self.assertIsInstance(df, pandas.DataFrame)
self.assertEqual(
mock_client_info.user_agent,
"gl-python/3.10.0",
)

def test_to_geodataframe_updates_user_agent(self):
pytest.importorskip("geopandas")
pyarrow = pytest.importorskip("pyarrow")
row_iterator = self._make_one_from_data(
(("name", "STRING"), ("geog", "GEOGRAPHY")),
(("foo", "Point(0 0)"),),
)
mock_client_info = mock.Mock(user_agent="test-agent")
row_iterator.client._connection = mock.Mock(_client_info=mock_client_info)
batch = pyarrow.RecordBatch.from_arrays(
[pyarrow.array(["foo"]), pyarrow.array(["Point(0 0)"])],
names=["name", "geog"],
)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=True,
),
mock.patch.object(row_iterator, "to_arrow", return_value=batch),
):
_ = row_iterator.to_geodataframe(create_bqstorage_client=False)

self.assertIn("pandas-gbq/", mock_client_info.user_agent)

def test_to_geodataframe_delegated_false_does_not_update_user_agent(self):
pytest.importorskip("geopandas")
pyarrow = pytest.importorskip("pyarrow")
row_iterator = self._make_one_from_data(
(("name", "STRING"), ("geog", "GEOGRAPHY")),
(("foo", "Point(0 0)"),),
)
mock_client_info = mock.Mock(user_agent="test-agent")
row_iterator.client._connection = mock.Mock(_client_info=mock_client_info)
batch = pyarrow.RecordBatch.from_arrays(
[pyarrow.array(["foo"]), pyarrow.array(["Point(0 0)"])],
names=["name", "geog"],
)

with (
mock.patch(
"google.cloud.bigquery._versions_helpers.PandasGBQVersions.is_delegation_supported",
new_callable=mock.PropertyMock,
return_value=False,
),
mock.patch.object(row_iterator, "to_arrow", return_value=batch),
):
_ = row_iterator.to_geodataframe(create_bqstorage_client=False)

self.assertEqual(mock_client_info.user_agent, "test-agent")


class TestPartitionRange(unittest.TestCase):
def _get_target_class(self):
Expand Down
Loading