ENH: Add user-agent string when constructing BigQuery and BigQuery Storage API clients. (#284)

* ENH: Add user-agent string when constructing BigQuery and BigQuery Storage API clients.

Since this was a relatively new addition to the BigQuery client library,
only populate the user-agent for BigQuery when recent-enough versions of
google-cloud-bigquery and google-api-core are installed.

* Add google-cloud-bigquery-storage to conda tests.

* Skip BigQuery Storage API test when package not available.

Use a smaller query to test the API.

* Make the BQ Storage API test query slightly smaller.

* Add fastavro to conda deps.

Author: tswast

ci/requirements-3.6-0.20.1.conda

@@ -1,7 +1,9 @@
-pydata-google-auth
-google-cloud-bigquery==1.9.0
-pytest
-pytest-cov
 codecov
 coverage
+fastavro
 flake8
+google-cloud-bigquery==1.9.0
+google-cloud-bigquery-storage==0.5.0
+pydata-google-auth
+pytest
+pytest-cov

ci/requirements-3.7-NIGHTLY.conda

@@ -1,5 +1,6 @@
 pydata-google-auth
-google-cloud-bigquery==1.10.0
+google-cloud-bigquery
+google-cloud-bigquery-storage
 pytest
 pytest-cov
 codecov

ci/requirements-3.7.pip

@@ -1,3 +1,3 @@
 pandas==0.24.0
-google-cloud-bigquery==1.9.0
+google-cloud-bigquery==1.12.0
 pydata-google-auth==0.1.2

docs/source/changelog.rst

@@ -16,6 +16,11 @@ Implementation changes
 - Use object dtype for ``STRING``, ``ARRAY``, and ``STRUCT`` columns when
   there are zero rows. (:issue:`285`)
 
+Internal changes
+~~~~~~~~~~~~~~~~
+
+- Populate ``user-agent`` with ``pandas`` version information. (:issue:`281`)
+
 .. _changelog-0.10.0:
 
 0.10.0 / 2019-04-05

pandas_gbq/gbq.py

@@ -8,16 +8,18 @@
 try:
     # The BigQuery Storage API client is an optional dependency. It is only
     # required when use_bqstorage_api=True.
-    from google.cloud import bigquery_storage_v1beta1
+    from google.cloud import bigquery_storage
 except ImportError:  # pragma: NO COVER
-    bigquery_storage_v1beta1 = None
+    bigquery_storage = None
 
 from pandas_gbq.exceptions import AccessDenied
 import pandas_gbq.schema
 
 logger = logging.getLogger(__name__)
 
 BIGQUERY_INSTALLED_VERSION = None
+BIGQUERY_CLIENT_INFO_VERSION = "1.12.0"
+HAS_CLIENT_INFO = False
 SHOW_VERBOSE_DEPRECATION = False
 SHOW_PRIVATE_KEY_DEPRECATION = False
 PRIVATE_KEY_DEPRECATION_MESSAGE = (
@@ -34,7 +36,7 @@
 
 
 def _check_google_client_version():
-    global BIGQUERY_INSTALLED_VERSION, SHOW_VERBOSE_DEPRECATION, SHOW_PRIVATE_KEY_DEPRECATION
+    global BIGQUERY_INSTALLED_VERSION, HAS_CLIENT_INFO, SHOW_VERBOSE_DEPRECATION, SHOW_PRIVATE_KEY_DEPRECATION
 
     try:
         import pkg_resources
@@ -44,10 +46,17 @@ def _check_google_client_version():
 
     # https://github.com/GoogleCloudPlatform/google-cloud-python/blob/master/bigquery/CHANGELOG.md
     bigquery_minimum_version = pkg_resources.parse_version("1.9.0")
+    bigquery_client_info_version = pkg_resources.parse_version(
+        BIGQUERY_CLIENT_INFO_VERSION
+    )
     BIGQUERY_INSTALLED_VERSION = pkg_resources.get_distribution(
         "google-cloud-bigquery"
     ).parsed_version
 
+    HAS_CLIENT_INFO = (
+        BIGQUERY_INSTALLED_VERSION >= bigquery_client_info_version
+    )
+
     if BIGQUERY_INSTALLED_VERSION < bigquery_minimum_version:
         raise ImportError(
             "pandas-gbq requires google-cloud-bigquery >= {0}, "
@@ -392,6 +401,29 @@ def sizeof_fmt(num, suffix="B"):
 
     def get_client(self):
         from google.cloud import bigquery
+        import pandas
+
+        try:
+            # This module was added in google-api-core 1.11.0.
+            # We don't have a hard requirement on that version, so only
+            # populate the client_info if available.
+            import google.api_core.client_info
+
+            client_info = google.api_core.client_info.ClientInfo(
+                user_agent="pandas-{}".format(pandas.__version__)
+            )
+        except ImportError:
+            client_info = None
+
+        # In addition to new enough version of google-api-core, a new enough
+        # version of google-cloud-bigquery is required to populate the
+        # client_info.
+        if HAS_CLIENT_INFO:
+            return bigquery.Client(
+                project=self.project_id,
+                credentials=self.credentials,
+                client_info=client_info,
+            )
 
         return bigquery.Client(
             project=self.project_id, credentials=self.credentials
@@ -751,14 +783,20 @@ def _make_bqstorage_client(use_bqstorage_api, credentials):
     if not use_bqstorage_api:
         return None
 
-    if bigquery_storage_v1beta1 is None:
+    if bigquery_storage is None:
         raise ImportError(
-            "Install the google-cloud-bigquery-storage and fastavro packages "
-            "to use the BigQuery Storage API."
+            "Install the google-cloud-bigquery-storage and fastavro/pyarrow "
+            "packages to use the BigQuery Storage API."
         )
 
-    return bigquery_storage_v1beta1.BigQueryStorageClient(
-        credentials=credentials
+    import google.api_core.gapic_v1.client_info
+    import pandas
+
+    client_info = google.api_core.gapic_v1.client_info.ClientInfo(
+        user_agent="pandas-{}".format(pandas.__version__)
+    )
+    return bigquery_storage.BigQueryStorageClient(
+        credentials=credentials, client_info=client_info
     )
 
 

tests/system/test_gbq.py

@@ -70,6 +70,16 @@ def random_dataset_id(bigquery_client):
         pass  # Not all tests actually create a dataset
 
 
+@pytest.fixture()
+def random_dataset(bigquery_client, random_dataset_id):
+    from google.cloud import bigquery
+
+    dataset_ref = bigquery_client.dataset(random_dataset_id)
+    dataset = bigquery.Dataset(dataset_ref)
+    bigquery_client.create_dataset(dataset)
+    return dataset
+
+
 @pytest.fixture()
 def tokyo_dataset(bigquery_client, random_dataset_id):
     from google.cloud import bigquery
@@ -894,31 +904,41 @@ def test_tokyo(self, tokyo_dataset, tokyo_table, project_id):
         assert df["max_year"][0] >= 2000
 
 
-@pytest.mark.skip(reason="large query for BQ Storage API tests")
-def test_read_gbq_w_bqstorage_api(credentials):
+@pytest.mark.slow(reason="Large query for BQ Storage API tests.")
+def test_read_gbq_w_bqstorage_api(credentials, random_dataset):
+    pytest.importorskip("google.cloud.bigquery_storage")
     df = gbq.read_gbq(
         """
         SELECT
-            dependency_name,
-            dependency_platform,
-            project_name,
-            project_id,
-            version_number,
-            version_id,
-            dependency_kind,
-            optional_dependency,
-            dependency_requirements,
-            dependency_project_id
-        FROM
-            `bigquery-public-data.libraries_io.dependencies`
-        WHERE
-            LOWER(dependency_platform) = 'npm'
-        LIMIT 2500000
+          total_amount,
+          passenger_count,
+          trip_distance
+        FROM `bigquery-public-data.new_york_taxi_trips.tlc_green_trips_2014`
+        -- Select non-null rows for no-copy conversion from Arrow to pandas.
+        WHERE total_amount IS NOT NULL
+          AND passenger_count IS NOT NULL
+          AND trip_distance IS NOT NULL
+        LIMIT 10000000
         """,
         use_bqstorage_api=True,
         credentials=credentials,
+        configuration={
+            "query": {
+                "destinationTable": {
+                    "projectId": random_dataset.project,
+                    "datasetId": random_dataset.dataset_id,
+                    "tableId": "".join(
+                        [
+                            "test_read_gbq_w_bqstorage_api_",
+                            str(uuid.uuid4()).replace("-", "_"),
+                        ]
+                    ),
+                },
+                "writeDisposition": "WRITE_TRUNCATE",
+            }
+        },
     )
-    assert len(df) == 2500000
+    assert len(df) == 10000000
 
 
 class TestToGBQIntegration(object):

tests/unit/conftest.py

@@ -15,12 +15,14 @@ def reset_context():
 
 @pytest.fixture(autouse=True)
 def mock_bigquery_client(monkeypatch):
-    from pandas_gbq import gbq
     from google.api_core.exceptions import NotFound
     import google.cloud.bigquery
     import google.cloud.bigquery.table
 
     mock_client = mock.create_autospec(google.cloud.bigquery.Client)
+    # Constructor returns the mock itself, so this mock can be treated as the
+    # constructor or the instance.
+    mock_client.return_value = mock_client
     mock_schema = [google.cloud.bigquery.SchemaField("_f0", "INTEGER")]
     # Mock out SELECT 1 query results.
     mock_query = mock.create_autospec(google.cloud.bigquery.QueryJob)
@@ -34,5 +36,6 @@ def mock_bigquery_client(monkeypatch):
     mock_client.query.return_value = mock_query
     # Mock table creation.
     mock_client.get_table.side_effect = NotFound("nope")
-    monkeypatch.setattr(gbq.GbqConnector, "get_client", lambda _: mock_client)
+    monkeypatch.setattr(google.cloud.bigquery, "Client", mock_client)
+    mock_client.reset_mock()
     return mock_client

tests/unit/test_auth.py

@@ -72,8 +72,6 @@ def mock_default_credentials(scopes=None, request=None):
         )
 
     monkeypatch.setattr(google.auth, "default", mock_default_credentials)
-    mock_client = mock.create_autospec(google.cloud.bigquery.Client)
-    monkeypatch.setattr(google.cloud.bigquery, "Client", mock_client)
 
     credentials, project = auth.get_credentials()
     assert project == "default-project"

tests/unit/test_gbq.py

@@ -5,6 +5,7 @@
 from unittest import mock
 
 import numpy
+import pandas
 from pandas import DataFrame
 import pandas.util.testing as tm
 import pkg_resources
@@ -22,6 +23,10 @@
 ).parsed_version
 
 
+def _make_connector(project_id="some-project", **kwargs):
+    return gbq.GbqConnector(project_id, **kwargs)
+
+
 @pytest.fixture
 def min_bq_version():
     import pkg_resources
@@ -99,7 +104,7 @@ def no_auth(monkeypatch):
         ("DATETIME", "datetime64[ns]"),
     ],
 )
-def test_should_return_bigquery_correctly_typed(type_, expected):
+def test__bqschema_to_nullsafe_dtypes(type_, expected):
     result = gbq._bqschema_to_nullsafe_dtypes(
         [dict(name="x", type=type_, mode="NULLABLE")]
     )
@@ -109,6 +114,35 @@ def test_should_return_bigquery_correctly_typed(type_, expected):
         assert result == {"x": expected}
 
 
+def test_GbqConnector_get_client_w_old_bq(monkeypatch, mock_bigquery_client):
+    gbq._test_google_api_imports()
+    connector = _make_connector()
+    monkeypatch.setattr(gbq, "HAS_CLIENT_INFO", False)
+
+    connector.get_client()
+
+    # No client_info argument.
+    mock_bigquery_client.assert_called_with(
+        credentials=mock.ANY, project=mock.ANY
+    )
+
+
+def test_GbqConnector_get_client_w_new_bq(mock_bigquery_client):
+    gbq._test_google_api_imports()
+    pytest.importorskip(
+        "google.cloud.bigquery", minversion=gbq.BIGQUERY_CLIENT_INFO_VERSION
+    )
+    pytest.importorskip("google.api_core.client_info")
+
+    connector = _make_connector()
+    connector.get_client()
+
+    _, kwargs = mock_bigquery_client.call_args
+    assert kwargs["client_info"].user_agent == "pandas-{}".format(
+        pandas.__version__
+    )
+
+
 def test_to_gbq_should_fail_if_invalid_table_name_passed():
     with pytest.raises(gbq.NotFoundException):
         gbq.to_gbq(DataFrame([[1]]), "invalid_table_name", project_id="1234")