diff --git a/CHANGELOG.md b/CHANGELOG.md index a8c2036bfc7..769375e21d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ ### Features - Added warning to nodes selector if nothing was matched ([#2115](https://github.com/fishtown-analytics/dbt/issues/2115), [#2343](https://github.com/fishtown-analytics/dbt/pull/2343)) - Suport column descriptions for BigQuery models ([#2335](https://github.com/fishtown-analytics/dbt/issues/2335), [#2402](https://github.com/fishtown-analytics/dbt/pull/2402)) - +- Added BigQuery option maximum_bytes_billed to set an upper limit for query costs ([#2346](https://github.com/fishtown-analytics/dbt/issues/2346), [#2427](https://github.com/fishtown-analytics/dbt/pull/2427)) ### Fixes - When tracking is disabled due to errors, do not reset the invocation ID ([#2398](https://github.com/fishtown-analytics/dbt/issues/2398), [#2400](https://github.com/fishtown-analytics/dbt/pull/2400)) @@ -29,6 +29,7 @@ Contributors: - [@azhard](https://github.com/azhard) ([#2413](https://github.com/fishtown-analytics/dbt/pull/2413), [#2422](https://github.com/fishtown-analytics/dbt/pull/2422)) - [@mikaelene](https://github.com/mikaelene) [#2414](https://github.com/fishtown-analytics/dbt/pull/2414) - [@raalsky](https://github.com/Raalsky) ([#2343](https://github.com/fishtown-analytics/dbt/pull/2343)) + - [@haukeduden](https://github.com/haukeduden) ([#2427](https://github.com/fishtown-analytics/dbt/pull/2427)) - [@alf-mindshift](https://github.com/alf-mindshift) ([docs#90](https://github.com/fishtown-analytics/dbt-docs/pull/90)) ## dbt 0.17.0b1 (May 5, 2020) diff --git a/plugins/bigquery/dbt/adapters/bigquery/connections.py b/plugins/bigquery/dbt/adapters/bigquery/connections.py index ac16068644a..ec07565ed0f 100644 --- a/plugins/bigquery/dbt/adapters/bigquery/connections.py +++ b/plugins/bigquery/dbt/adapters/bigquery/connections.py @@ -43,6 +43,7 @@ class BigQueryCredentials(Credentials): location: Optional[str] = None priority: Optional[Priority] = None retries: Optional[int] = 1 + maximum_bytes_billed: Optional[int] = None _ALIASES = { 'project': 'database', 'dataset': 'schema', @@ -54,7 +55,7 @@ def type(self): def _connection_keys(self): return ('method', 'database', 'schema', 'location', 'priority', - 'timeout_seconds') + 'timeout_seconds', 'maximum_bytes_billed') class BigQueryConnectionManager(BaseConnectionManager): @@ -216,6 +217,10 @@ def raw_execute(self, sql, fetch=False): job_params[ 'priority'] = google.cloud.bigquery.QueryPriority.INTERACTIVE + maximum_bytes_billed = conn.credentials.maximum_bytes_billed + if maximum_bytes_billed is not None and maximum_bytes_billed != 0: + job_params['maximum_bytes_billed'] = maximum_bytes_billed + def fn(): return self._query_and_results(client, sql, conn, job_params) diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py index 456b6d79ded..4a8ac7ae303 100644 --- a/test/unit/test_bigquery_adapter.py +++ b/test/unit/test_bigquery_adapter.py @@ -58,6 +58,7 @@ def setUp(self): 'threads': 1, 'location': 'Luna Station', 'priority': 'batch', + 'maximum_bytes_billed': 0, }, }, 'target': 'oauth', @@ -147,6 +148,21 @@ def test_acquire_connection_priority(self, mock_open_connection): connection.handle mock_open_connection.assert_called_once() + @patch('dbt.adapters.bigquery.BigQueryConnectionManager.open', return_value=_bq_conn()) + def test_acquire_connection_maximum_bytes_billed(self, mock_open_connection): + adapter = self.get_adapter('loc') + try: + connection = adapter.acquire_connection('dummy') + self.assertEqual(connection.type, 'bigquery') + self.assertEqual(connection.credentials.maximum_bytes_billed, 0) + + except dbt.exceptions.ValidationException as e: + self.fail('got ValidationException: {}'.format(str(e))) + + mock_open_connection.assert_not_called() + connection.handle + mock_open_connection.assert_called_once() + def test_cancel_open_connections_empty(self): adapter = self.get_adapter('oauth') self.assertEqual(adapter.cancel_open_connections(), None)