fix: correct user agent header to track calls as veneer instead of gapic (#909)

* fix: correct user agent header to track calls as veneer instead of gapic

* add unit test

* move where client info is changed

* correction

* change header for v1beta2, too

Author: Linchin

google/cloud/bigquery_storage_v1/client.py

@@ -21,8 +21,10 @@
 
 from __future__ import absolute_import
 
+from google.api_core import gapic_v1
 import google.api_core.gapic_v1.method
 
+from google.cloud.bigquery_storage_v1 import gapic_version as package_version
 from google.cloud.bigquery_storage_v1 import reader
 from google.cloud.bigquery_storage_v1.services import big_query_read, big_query_write
 
@@ -31,13 +33,22 @@
     "https://www.googleapis.com/auth/cloud-platform",
 )
 
+VENEER_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
+    client_library_version=package_version.__version__
+)
+
 
 class BigQueryReadClient(big_query_read.BigQueryReadClient):
     """Client for interacting with BigQuery Storage API.
 
     The BigQuery storage API can be used to read data stored in BigQuery.
     """
 
+    def __init__(self, **kwargs):
+        if "client_info" not in kwargs:
+            kwargs["client_info"] = VENEER_CLIENT_INFO
+        super().__init__(**kwargs)
+
     def read_rows(
         self,
         name,
@@ -140,3 +151,8 @@ def read_rows(
 
 class BigQueryWriteClient(big_query_write.BigQueryWriteClient):
     __doc__ = big_query_write.BigQueryWriteClient.__doc__
+
+    def __init__(self, **kwargs):
+        if "client_info" not in kwargs:
+            kwargs["client_info"] = VENEER_CLIENT_INFO
+        super().__init__(**kwargs)

google/cloud/bigquery_storage_v1beta2/client.py

@@ -19,9 +19,11 @@
 This is the base from which all interactions with the API occur.
 """
 
+from google.api_core import gapic_v1
 import google.api_core.gapic_v1.method
 import google.api_core.retry
 
+from google.cloud.bigquery_storage_v1 import gapic_version as package_version
 from google.cloud.bigquery_storage_v1 import reader
 from google.cloud.bigquery_storage_v1beta2.services import (
     big_query_read,
@@ -33,13 +35,22 @@
     "https://www.googleapis.com/auth/cloud-platform",
 )
 
+VENEER_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
+    client_library_version=package_version.__version__
+)
+
 
 class BigQueryReadClient(big_query_read.BigQueryReadClient):
     """Client for interacting with BigQuery Storage API.
 
     The BigQuery storage API can be used to read data stored in BigQuery.
     """
 
+    def __init__(self, **kwargs):
+        if "client_info" not in kwargs:
+            kwargs["client_info"] = VENEER_CLIENT_INFO
+        super().__init__(**kwargs)
+
     def read_rows(
         self,
         name,
@@ -142,3 +153,8 @@ def read_rows(
 
 class BigQueryWriteClient(big_query_write.BigQueryWriteClient):
     __doc__ = big_query_write.BigQueryWriteClient.__doc__
+
+    def __init__(self, **kwargs):
+        if "client_info" not in kwargs:
+            kwargs["client_info"] = VENEER_CLIENT_INFO
+        super().__init__(**kwargs)

tests/unit/test_client_v1.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 from unittest import mock
 
 from google.api_core.gapic_v1 import client_info
+from google.auth import credentials
 import pytest
 
 from google.cloud.bigquery_storage import types
@@ -115,3 +117,31 @@ def test_read_rows(mock_transport, client_under_test):
     mock_transport.create_read_session.read_rows(
         expected_request, metadata=mock.ANY, timeout=mock.ANY
     )
+
+
+@pytest.mark.parametrize(
+    "module_under_test",
+    ["google.cloud.bigquery_storage_v1", "google.cloud.bigquery_storage_v1beta2"],
+)
+def test_init_default_client_info(module_under_test):
+    from google.api_core.gapic_v1.client_info import METRICS_METADATA_KEY
+
+    mut = importlib.import_module(module_under_test)
+
+    creds = mock.Mock(spec=credentials.Credentials)
+    client = mut.BigQueryWriteClient(credentials=creds)
+
+    installed_version = mut.__version__
+    expected_client_info = f"gccl/{installed_version}"
+
+    for wrapped_method in client.transport._wrapped_methods.values():
+        user_agent = next(
+            (
+                header_value
+                for header, header_value in wrapped_method._metadata
+                if header == METRICS_METADATA_KEY
+            ),
+            None,
+        )
+        assert user_agent is not None
+        assert expected_client_info in user_agent