diff --git a/CHANGELOG.md b/CHANGELOG.md index f710998d614..245279361b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ### Features - Add a "docs" field to models, with a "show" subfield ([#1671](https://github.com/fishtown-analytics/dbt/issues/1671), [#2107](https://github.com/fishtown-analytics/dbt/pull/2107)) +- Add a dbt-{dbt_version} user agent field to the bigquery connector ([#2121](https://github.com/fishtown-analytics/dbt/issues/2121), [#2146](https://github.com/fishtown-analytics/dbt/pull/2146)) ### Fixes - Fix issue where dbt did not give an error in the presence of duplicate doc names ([#2054](https://github.com/fishtown-analytics/dbt/issues/2054), [#2080](https://github.com/fishtown-analytics/dbt/pull/2080)) diff --git a/plugins/bigquery/dbt/adapters/bigquery/connections.py b/plugins/bigquery/dbt/adapters/bigquery/connections.py index 35146bee24f..82fd74ead75 100644 --- a/plugins/bigquery/dbt/adapters/bigquery/connections.py +++ b/plugins/bigquery/dbt/adapters/bigquery/connections.py @@ -3,16 +3,18 @@ from typing import Optional, Any, Dict import google.auth -import google.api_core -import google.oauth2 -import google.cloud.exceptions import google.cloud.bigquery -from google.api_core import retry +import google.cloud.exceptions +from google.api_core import retry, client_info +from google.oauth2 import service_account -import dbt.clients.agate_helper -import dbt.exceptions +from dbt.clients import agate_helper, gcloud +from dbt.exceptions import ( + FailedToConnectException, RuntimeException, DatabaseException +) from dbt.adapters.base import BaseConnectionManager, Credentials from dbt.logger import GLOBAL_LOGGER as logger +from dbt.version import __version__ as dbt_version from hologram.helpers import StrEnum @@ -70,7 +72,7 @@ def handle_error(cls, error, message, sql): error_msg = "\n".join( [item['message'] for item in error.errors]) - raise dbt.exceptions.DatabaseException(error_msg) from error + raise DatabaseException(error_msg) from error def clear_transaction(self): pass @@ -91,12 +93,12 @@ def exception_handler(self, sql): except Exception as e: logger.debug("Unhandled error while running:\n{}".format(sql)) logger.debug(e) - if isinstance(e, dbt.exceptions.RuntimeException): + if isinstance(e, RuntimeException): # during a sql query, an internal to dbt exception was raised. # this sounds a lot like a signal handler and probably has # useful information, so raise it without modification. raise - raise dbt.exceptions.RuntimeException(str(e)) from e + raise RuntimeException(str(e)) from e def cancel_open(self) -> None: pass @@ -116,7 +118,7 @@ def commit(self): @classmethod def get_bigquery_credentials(cls, profile_credentials): method = profile_credentials.method - creds = google.oauth2.service_account.Credentials + creds = service_account.Credentials if method == BigQueryConnectionMethod.OAUTH: credentials, project_id = google.auth.default(scopes=cls.SCOPE) @@ -131,15 +133,21 @@ def get_bigquery_credentials(cls, profile_credentials): return creds.from_service_account_info(details, scopes=cls.SCOPE) error = ('Invalid `method` in profile: "{}"'.format(method)) - raise dbt.exceptions.FailedToConnectException(error) + raise FailedToConnectException(error) @classmethod def get_bigquery_client(cls, profile_credentials): database = profile_credentials.database creds = cls.get_bigquery_credentials(profile_credentials) location = getattr(profile_credentials, 'location', None) - return google.cloud.bigquery.Client(database, creds, - location=location) + + info = client_info.ClientInfo(user_agent=f'dbt-{dbt_version}') + return google.cloud.bigquery.Client( + database, + creds, + location=location, + client_info=info, + ) @classmethod def open(cls, connection): @@ -152,7 +160,7 @@ def open(cls, connection): except google.auth.exceptions.DefaultCredentialsError: logger.info("Please log into GCP to continue") - dbt.clients.gcloud.setup_default_credentials() + gcloud.setup_default_credentials() handle = cls.get_bigquery_client(connection.credentials) @@ -164,7 +172,7 @@ def open(cls, connection): connection.handle = None connection.state = 'fail' - raise dbt.exceptions.FailedToConnectException(str(e)) + raise FailedToConnectException(str(e)) connection.handle = handle connection.state = 'open' @@ -186,8 +194,7 @@ def get_retries(cls, conn) -> int: @classmethod def get_table_from_response(cls, resp): column_names = [field.name for field in resp.schema] - return dbt.clients.agate_helper.table_from_data_flat(resp, - column_names) + return agate_helper.table_from_data_flat(resp, column_names) def raw_execute(self, sql, fetch=False): conn = self.get_thread_connection() @@ -219,7 +226,7 @@ def execute(self, sql, auto_begin=False, fetch=None): if fetch: res = self.get_table_from_response(iterator) else: - res = dbt.clients.agate_helper.empty_table() + res = agate_helper.empty_table() if query_job.statement_type == 'CREATE_VIEW': status = 'CREATE VIEW' diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py index 05e738c9a09..adaf20bcedb 100644 --- a/test/unit/test_bigquery_adapter.py +++ b/test/unit/test_bigquery_adapter.py @@ -1,3 +1,4 @@ +import re import unittest from contextlib import contextmanager from unittest.mock import patch, MagicMock, Mock @@ -173,7 +174,16 @@ def test_location_value(self, mock_bq, mock_auth_default): mock_client.assert_not_called() connection.handle mock_client.assert_called_once_with('dbt-unit-000000', creds, - location='Luna Station') + location='Luna Station', + client_info=HasUserAgent()) + + +class HasUserAgent: + PAT = re.compile(r'dbt-\d+\.\d+\.\d+[a-zA-Z]+\d+') + + def __eq__(self, other): + compare = getattr(other, 'user_agent', '') + return bool(self.PAT.match(compare)) class TestConnectionNamePassthrough(BaseTestBigQueryAdapter):