From 4e3f86e02a98e7c0ae9b17c7e14ee38de08314e3 Mon Sep 17 00:00:00 2001 From: Jack Date: Fri, 1 Nov 2024 20:05:23 +0000 Subject: [PATCH] Feature/add quota project option (#1345) * add ability to specify quota project in profile --------- Co-authored-by: Mike Alfare <13974384+mikealfare@users.noreply.github.com> --- .../unreleased/Features-20240911-234859.yaml | 6 +++++ dbt/adapters/bigquery/connections.py | 6 ++++- tests/functional/test_quota_project.py | 27 +++++++++++++++++++ tests/unit/test_bigquery_adapter.py | 5 +++- 4 files changed, 42 insertions(+), 2 deletions(-) create mode 100644 .changes/unreleased/Features-20240911-234859.yaml create mode 100644 tests/functional/test_quota_project.py diff --git a/.changes/unreleased/Features-20240911-234859.yaml b/.changes/unreleased/Features-20240911-234859.yaml new file mode 100644 index 000000000..5351c3315 --- /dev/null +++ b/.changes/unreleased/Features-20240911-234859.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Adds the ability to set optional `quota_project` in profile +time: 2024-09-11T23:48:59.767649+01:00 +custom: + Author: jcarpenter12 + Issue: 1343 1344 diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index d3eee3ef3..58b3dbe41 100644 --- a/dbt/adapters/bigquery/connections.py +++ b/dbt/adapters/bigquery/connections.py @@ -17,7 +17,7 @@ import google.auth.exceptions import google.cloud.bigquery import google.cloud.exceptions -from google.api_core import retry, client_info +from google.api_core import retry, client_info, client_options from google.auth import impersonated_credentials from google.oauth2 import ( credentials as GoogleCredentials, @@ -125,6 +125,7 @@ class BigQueryCredentials(Credentials): database: Optional[str] = None schema: Optional[str] = None execution_project: Optional[str] = None + quota_project: Optional[str] = None location: Optional[str] = None priority: Optional[Priority] = None maximum_bytes_billed: Optional[int] = None @@ -408,14 +409,17 @@ def get_credentials(cls, profile_credentials): def get_bigquery_client(cls, profile_credentials): creds = cls.get_credentials(profile_credentials) execution_project = profile_credentials.execution_project + quota_project = profile_credentials.quota_project location = getattr(profile_credentials, "location", None) info = client_info.ClientInfo(user_agent=f"dbt-bigquery-{dbt_version.version}") + options = client_options.ClientOptions(quota_project_id=quota_project) return google.cloud.bigquery.Client( execution_project, creds, location=location, client_info=info, + client_options=options, ) @classmethod diff --git a/tests/functional/test_quota_project.py b/tests/functional/test_quota_project.py new file mode 100644 index 000000000..0b4bb90c4 --- /dev/null +++ b/tests/functional/test_quota_project.py @@ -0,0 +1,27 @@ +import os + +import pytest + +from dbt.tests.util import run_dbt + +_QUOTA_PROJECT = os.getenv("BIGQUERY_TEST_ALT_DATABASE") + + +class TestNoQuotaProject: + def test_no_quota_project(self, project): + results = run_dbt() + for result in results: + assert None == result.adapter_response["quota_project"] + + +class TestQuotaProjectOption: + @pytest.fixture(scope="class") + def profiles_config_update(self, dbt_profile_target): + outputs = {"default": dbt_profile_target} + outputs["default"]["quota_project"] = _QUOTA_PROJECT + yield + + def test_quota_project_option(self, project): + results = run_dbt() + for result in results: + assert _QUOTA_PROJECT == result.adapter_response["quota_project"] diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index a922525fd..da499d266 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -386,15 +386,17 @@ def test_cancel_open_connections_single(self): adapter.connections.thread_connections.update({key: master, 1: model}) self.assertEqual(len(list(adapter.cancel_open_connections())), 1) + @patch("dbt.adapters.bigquery.impl.google.api_core.client_options.ClientOptions") @patch("dbt.adapters.bigquery.impl.google.auth.default") @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery") - def test_location_user_agent(self, mock_bq, mock_auth_default): + def test_location_user_agent(self, mock_bq, mock_auth_default, MockClientOptions): creds = MagicMock() mock_auth_default.return_value = (creds, MagicMock()) adapter = self.get_adapter("loc") connection = adapter.acquire_connection("dummy") mock_client = mock_bq.Client + mock_client_options = MockClientOptions.return_value mock_client.assert_not_called() connection.handle @@ -403,6 +405,7 @@ def test_location_user_agent(self, mock_bq, mock_auth_default): creds, location="Luna Station", client_info=HasUserAgent(), + client_options=mock_client_options, )