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