diff --git a/google/cloud/bigquery_storage_v1/client.py b/google/cloud/bigquery_storage_v1/client.py index c891a1b8..330c4a13 100644 --- a/google/cloud/bigquery_storage_v1/client.py +++ b/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,6 +33,10 @@ "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. @@ -38,6 +44,11 @@ class BigQueryReadClient(big_query_read.BigQueryReadClient): 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) diff --git a/google/cloud/bigquery_storage_v1beta2/client.py b/google/cloud/bigquery_storage_v1beta2/client.py index a02f7651..323b5941 100644 --- a/google/cloud/bigquery_storage_v1beta2/client.py +++ b/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,6 +35,10 @@ "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. @@ -40,6 +46,11 @@ class BigQueryReadClient(big_query_read.BigQueryReadClient): 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) diff --git a/tests/unit/test_client_v1.py b/tests/unit/test_client_v1.py index 0cb7466e..f214b877 100644 --- a/tests/unit/test_client_v1.py +++ b/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