Skip to content

Commit

Permalink
first try
Browse files Browse the repository at this point in the history
  • Loading branch information
bbhoss committed Aug 1, 2020
1 parent 1fb8c9c commit 88529d5
Showing 1 changed file with 26 additions and 1 deletion.
27 changes: 26 additions & 1 deletion plugins/bigquery/dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,11 @@
from typing import Optional, Any, Dict

import google.auth
import google.auth.exceptions
import google.cloud.bigquery
import google.cloud.exceptions
from google.api_core import retry, client_info
from google.auth import impersonated_credentials
from google.oauth2 import service_account

from dbt.utils import format_bytes, format_rows_number
Expand Down Expand Up @@ -45,6 +47,7 @@ class BigQueryCredentials(Credentials):
priority: Optional[Priority] = None
retries: Optional[int] = 1
maximum_bytes_billed: Optional[int] = None
impersonate_service_account: Optional[str] = None
_ALIASES = {
'project': 'database',
'dataset': 'schema',
Expand Down Expand Up @@ -92,6 +95,14 @@ def exception_handler(self, sql):
message = "Access denied while running query"
self.handle_error(e, message)

except google.auth.exceptions.RefreshError:
message = "Unable to generate access token, if you're using " \
"impersonate_service_account, make sure your " \
'initial account has the "roles/' \
'iam.serviceAccountTokenCreator" role on the ' \
'account you are trying to impersonate.'
raise RuntimeException(message)

except Exception as e:
logger.debug("Unhandled error while running:\n{}".format(sql))
logger.debug(e)
Expand Down Expand Up @@ -142,10 +153,24 @@ def get_bigquery_credentials(cls, profile_credentials):
error = ('Invalid `method` in profile: "{}"'.format(method))
raise FailedToConnectException(error)

@classmethod
def get_impersonated_bigquery_credentials(cls, profile_credentials):
source_credentials = cls.get_bigquery_credentials(profile_credentials)
return impersonated_credentials.Credentials(
source_credentials=source_credentials,
target_principal=profile_credentials.impersonate_service_account,
target_scopes=list(cls.SCOPE),
lifetime=profile_credentials.timeout_seconds,
)

@classmethod
def get_bigquery_client(cls, profile_credentials):
if profile_credentials.impersonate_service_account:
creds =\
cls.get_impersonated_bigquery_credentials(profile_credentials)
else:
creds = cls.get_bigquery_credentials(profile_credentials)
database = profile_credentials.database
creds = cls.get_bigquery_credentials(profile_credentials)
location = getattr(profile_credentials, 'location', None)

info = client_info.ClientInfo(user_agent=f'dbt-{dbt_version}')
Expand Down

0 comments on commit 88529d5

Please sign in to comment.