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

packages/google-cloud-bigquery/google/cloud/bigquery/table.py

@@ -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:
@@ -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
 
@@ -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
 
@@ -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()
+
         record_batch = self.to_arrow(
             progress_bar_type=progress_bar_type,
             bqstorage_client=bqstorage_client,
@@ -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:

packages/google-cloud-bigquery/tests/unit/test_table.py

@@ -16,6 +16,7 @@
 import datetime
 import logging
 import re
+import sys
 import time
 import types
 import unittest
@@ -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):