Merge pull request #2677 from bbhoss/bq_impersonate

Add support for impersonating a service account with BigQuery

CHANGELOG.md

@@ -1,12 +1,20 @@
 ## dbt 0.18.0 (Release TBD)
 
+
+### Features
+- Add support for impersonating a service account using `impersonate_service_account` in the BigQuery profile configuration ([#2677](https://github.com/fishtown-analytics/dbt/issues/2677)) ([docs](https://docs.getdbt.com/reference/warehouse-profiles/bigquery-profile#service-account-impersonation))
+
+
 ### Breaking changes
 - `adapter_macro` is no longer a macro, instead it is a builtin context method. Any custom macros that intercepted it by going through `context['dbt']` will need to instead access it via `context['builtins']` ([#2302](https://github.com/fishtown-analytics/dbt/issues/2302), [#2673](https://github.com/fishtown-analytics/dbt/pull/2673))
 - `adapter_macro` is now deprecated. Use `adapter.dispatch` instead.
 
 ### Features
 - Added a `dispatch` method to the context adapter and deprecated `adapter_macro`. ([#2302](https://github.com/fishtown-analytics/dbt/issues/2302), [#2679](https://github.com/fishtown-analytics/dbt/pull/2679))
 
+Contributors:
+- [@bbhoss](https://github.com/bbhoss) ([#2677](https://github.com/fishtown-analytics/dbt/pull/2677))
+
 ## dbt 0.18.0b2 (July 30, 2020)
 
 ### Features

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

@@ -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
@@ -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',
@@ -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)
@@ -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}')

test/unit/test_bigquery_adapter.py

@@ -61,6 +61,14 @@ def setUp(self):
                     'priority': 'batch',
                     'maximum_bytes_billed': 0,
                 },
+                'impersonate': {
+                    'type': 'bigquery',
+                    'method': 'oauth',
+                    'project': 'dbt-unit-000000',
+                    'schema': 'dummy_schema',
+                    'threads': 1,
+                    'impersonate_service_account': '[email protected]'
+                },
             },
             'target': 'oauth',
         }
@@ -134,6 +142,23 @@ 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_impersonated_service_account_validations(self, mock_open_connection):
+        adapter = self.get_adapter('impersonate')
+        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_priority(self, mock_open_connection):
         adapter = self.get_adapter('loc')