Skip to content

Commit

Permalink
add a user agent with version to the bigquery client info
Browse files Browse the repository at this point in the history
Updated unit tests
  • Loading branch information
Jacob Beck committed Feb 18, 2020
1 parent 542a653 commit b5f8800
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 19 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
43 changes: 25 additions & 18 deletions plugins/bigquery/dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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):
Expand All @@ -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)

Expand All @@ -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'
Expand All @@ -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()
Expand Down Expand Up @@ -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'
Expand Down
12 changes: 11 additions & 1 deletion test/unit/test_bigquery_adapter.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import re
import unittest
from contextlib import contextmanager
from unittest.mock import patch, MagicMock, Mock
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit b5f8800

Please sign in to comment.