Support BigQuery OAuth using a refresh token and client secrets

CHANGELOG.md

@@ -13,6 +13,7 @@
 - Save manifest at the same time we save the run_results at the end of a run ([#2765](https://github.com/fishtown-analytics/dbt/issues/2765), [#2799](https://github.com/fishtown-analytics/dbt/pull/2799))
 - Added dbt_invocation_id for each BigQuery job to enable performance analysis ([#2808](https://github.com/fishtown-analytics/dbt/issues/2808), [#2809](https://github.com/fishtown-analytics/dbt/pull/2809))
 - Save cli and rpc arguments in run_results.json ([#2510](https://github.com/fishtown-analytics/dbt/issues/2510), [#2813](https://github.com/fishtown-analytics/dbt/pull/2813))
+- Added support for BigQuery connections using refresh tokens ([#2344](https://github.com/fishtown-analytics/dbt/issues/2344), [#2805](https://github.com/fishtown-analytics/dbt/pull/2805))
 
 ### Under the hood
 - Added strategy-specific validation to improve the relevancy of compilation errors for the `timestamp` and `check` snapshot strategies. (([#2787](https://github.com/fishtown-analytics/dbt/issues/2787), [#2791](https://github.com/fishtown-analytics/dbt/pull/2791))

plugins/bigquery/dbt/adapters/bigquery/connections.py

@@ -9,7 +9,10 @@
 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 google.oauth2 import (
+    credentials as GoogleCredentials,
+    service_account as GoogleServiceAccountCredentials
+)
 
 from dbt.utils import format_bytes, format_rows_number
 from dbt.clients import agate_helper, gcloud
@@ -51,19 +54,30 @@ class BigQueryConnectionMethod(StrEnum):
     OAUTH = 'oauth'
     SERVICE_ACCOUNT = 'service-account'
     SERVICE_ACCOUNT_JSON = 'service-account-json'
+    OAUTH_SECRETS = 'oauth-secrets'
 
 
 @dataclass
 class BigQueryCredentials(Credentials):
     method: BigQueryConnectionMethod
-    keyfile: Optional[str] = None
-    keyfile_json: Optional[Dict[str, Any]] = None
     timeout_seconds: Optional[int] = 300
     location: Optional[str] = None
     priority: Optional[Priority] = None
     retries: Optional[int] = 1
     maximum_bytes_billed: Optional[int] = None
     impersonate_service_account: Optional[str] = None
+
+    # Keyfile json creds
+    keyfile: Optional[str] = None
+    keyfile_json: Optional[Dict[str, Any]] = None
+
+    # oauth-secrets
+    token: Optional[str] = None
+    refresh_token: Optional[str] = None
+    client_id: Optional[str] = None
+    client_secret: Optional[str] = None
+    token_uri: Optional[str] = None
+
     _ALIASES = {
         'project': 'database',
         'dataset': 'schema',
@@ -111,12 +125,13 @@ def exception_handler(self, sql):
             message = "Access denied while running query"
             self.handle_error(e, message)
 
-        except google.auth.exceptions.RefreshError:
+        except google.auth.exceptions.RefreshError as e:
             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.'
+                      'account you are trying to impersonate.\n\n' \
+                      f'{str(e)}'
             raise RuntimeException(message)
 
         except Exception as e:
@@ -152,7 +167,7 @@ def commit(self):
     @classmethod
     def get_bigquery_credentials(cls, profile_credentials):
         method = profile_credentials.method
-        creds = service_account.Credentials
+        creds = GoogleServiceAccountCredentials.Credentials
 
         if method == BigQueryConnectionMethod.OAUTH:
             credentials, project_id = google.auth.default(scopes=cls.SCOPE)
@@ -166,6 +181,16 @@ def get_bigquery_credentials(cls, profile_credentials):
             details = profile_credentials.keyfile_json
             return creds.from_service_account_info(details, scopes=cls.SCOPE)
 
+        elif method == BigQueryConnectionMethod.OAUTH_SECRETS:
+            return GoogleCredentials.Credentials(
+                token=profile_credentials.token,
+                refresh_token=profile_credentials.refresh_token,
+                client_id=profile_credentials.client_id,
+                client_secret=profile_credentials.client_secret,
+                token_uri=profile_credentials.token_uri,
+                scopes=cls.SCOPE
+            )
+
         error = ('Invalid `method` in profile: "{}"'.format(method))
         raise FailedToConnectException(error)
 

test/unit/test_bigquery_adapter.py

@@ -71,6 +71,31 @@ def setUp(self):
                     'threads': 1,
                     'impersonate_service_account': '[email protected]'
                 },
+                'oauth-credentials-token': {
+                    'type': 'bigquery',
+                    'method': 'oauth-secrets',
+                    'token': 'abc',
+                    'project': 'dbt-unit-000000',
+                    'schema': 'dummy_schema',
+                    'threads': 1,
+                    'location': 'Luna Station',
+                    'priority': 'batch',
+                    'maximum_bytes_billed': 0,
+                },
+                'oauth-credentials': {
+                    'type': 'bigquery',
+                    'method': 'oauth-secrets',
+                    'client_id': 'abc',
+                    'client_secret': 'def',
+                    'refresh_token': 'ghi',
+                    'token_uri': 'jkl',
+                    'project': 'dbt-unit-000000',
+                    'schema': 'dummy_schema',
+                    'threads': 1,
+                    'location': 'Luna Station',
+                    'priority': 'batch',
+                    'maximum_bytes_billed': 0,
+                },
             },
             'target': 'oauth',
         }
@@ -145,6 +170,40 @@ def test_acquire_connection_service_account_validations(self, mock_open_connecti
         connection.handle
         mock_open_connection.assert_called_once()
 
+    @patch('dbt.adapters.bigquery.BigQueryConnectionManager.open', return_value=_bq_conn())
+    def test_acquire_connection_oauth_token_validations(self, mock_open_connection):
+        adapter = self.get_adapter('oauth-credentials-token')
+        try:
+            connection = adapter.acquire_connection('dummy')
+            self.assertEqual(connection.type, 'bigquery')
+
+        except dbt.exceptions.ValidationException as e:
+            self.fail('got ValidationException: {}'.format(str(e)))
+
+        except BaseException as e:
+            raise
+
+        mock_open_connection.assert_not_called()
+        connection.handle
+        mock_open_connection.assert_called_once()
+
+    @patch('dbt.adapters.bigquery.BigQueryConnectionManager.open', return_value=_bq_conn())
+    def test_acquire_connection_oauth_credentials_validations(self, mock_open_connection):
+        adapter = self.get_adapter('oauth-credentials')
+        try:
+            connection = adapter.acquire_connection('dummy')
+            self.assertEqual(connection.type, 'bigquery')
+
+        except dbt.exceptions.ValidationException as e:
+            self.fail('got ValidationException: {}'.format(str(e)))
+
+        except BaseException as e:
+            raise
+
+        mock_open_connection.assert_not_called()
+        connection.handle
+        mock_open_connection.assert_called_once()
+
     @patch('dbt.adapters.bigquery.BigQueryConnectionManager.open', return_value=_bq_conn())
     def test_acquire_connection_impersonated_service_account_validations(self, mock_open_connection):
         adapter = self.get_adapter('impersonate')