Feature/add quota project option (#1345)

* add ability to specify quota project in profile

---------

Co-authored-by: Mike Alfare <[email protected]>

Author: jcarpenter12

.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

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

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"]

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,
         )