From 88529d5c252c529954dcc605baa5686a8ce162aa Mon Sep 17 00:00:00 2001 From: Preston Marshall Date: Sat, 1 Aug 2020 15:12:00 -0400 Subject: [PATCH 1/5] first try --- .../dbt/adapters/bigquery/connections.py | 27 ++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/plugins/bigquery/dbt/adapters/bigquery/connections.py b/plugins/bigquery/dbt/adapters/bigquery/connections.py index d18c4ec9210..9ed789e62fb 100644 --- a/plugins/bigquery/dbt/adapters/bigquery/connections.py +++ b/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}') From c95a6792e5509904c3bc42d0c3d7f6c8ace11ca3 Mon Sep 17 00:00:00 2001 From: Preston Marshall Date: Sat, 1 Aug 2020 15:40:48 -0400 Subject: [PATCH 2/5] add a test --- test/unit/test_bigquery_adapter.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py index 71c3c95d3b8..c0e9198850a 100644 --- a/test/unit/test_bigquery_adapter.py +++ b/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': 'dummyaccount@dbt.iam.gserviceaccount.com' + }, }, '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') From 7a9fc7ef1298150f91a0a422ab3ee696b3bed3bd Mon Sep 17 00:00:00 2001 From: Preston Marshall Date: Mon, 3 Aug 2020 12:55:42 -0400 Subject: [PATCH 3/5] update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4a24580e6f..f4f0195664b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Features - Added `--defer` and `--state` flags to `dbt run`, to defer to a previously generated manifest for unselected nodes in a run. ([#2527](https://github.com/fishtown-analytics/dbt/issues/2527), [#2656](https://github.com/fishtown-analytics/dbt/pull/2656)) +- 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 From 43d5dfcb71a5f35c539456d2fb72120dfce39060 Mon Sep 17 00:00:00 2001 From: Preston Marshall Date: Mon, 3 Aug 2020 12:57:30 -0400 Subject: [PATCH 4/5] move changelog additions to next --- CHANGELOG.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f4f0195664b..7b23e347e84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,13 @@ ## 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)) + ## dbt 0.18.0b2 (July 30, 2020) ### Features - Added `--defer` and `--state` flags to `dbt run`, to defer to a previously generated manifest for unselected nodes in a run. ([#2527](https://github.com/fishtown-analytics/dbt/issues/2527), [#2656](https://github.com/fishtown-analytics/dbt/pull/2656)) -- 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 From df8ccc04eb297442e6efc71b3155dadbc06ecff1 Mon Sep 17 00:00:00 2001 From: Preston Marshall Date: Tue, 4 Aug 2020 10:49:35 -0400 Subject: [PATCH 5/5] add contributors section --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 67b1e006e1f..e71dc9bd069 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,9 +4,13 @@ ### 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)) +Contributors: +- [@bbhoss](https://github.com/bbhoss) ([#2677](https://github.com/fishtown-analytics/dbt/pull/2677)) + ## dbt 0.18.0b2 (July 30, 2020) ### Features