diff --git a/dbt-bigquery/.changes/unreleased/Features-20250501-081844.yaml b/dbt-bigquery/.changes/unreleased/Features-20250501-081844.yaml new file mode 100644 index 000000000..a3e14b56f --- /dev/null +++ b/dbt-bigquery/.changes/unreleased/Features-20250501-081844.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Make Bigquery `api_endpoint` config option available for dbt client +time: 2025-05-01T08:18:44.483014492Z +custom: + Author: OTooleMichael + Issue: https://github.com/dbt-labs/dbt-bigquery/pull/1017 diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py index 90c250dc5..2d9121791 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py @@ -61,7 +61,9 @@ def _create_bigquery_client(credentials: BigQueryCredentials) -> BigQueryClient: create_google_credentials(credentials), location=getattr(credentials, "location", None), client_info=ClientInfo(user_agent=f"dbt-bigquery-{dbt_version.version}"), - client_options=ClientOptions(quota_project_id=credentials.quota_project), + client_options=ClientOptions( + quota_project_id=credentials.quota_project, api_endpoint=credentials.api_endpoint + ), ) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/credentials.py b/dbt-bigquery/src/dbt/adapters/bigquery/credentials.py index 437337f53..28aee2104 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/credentials.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/credentials.py @@ -55,6 +55,7 @@ class BigQueryCredentials(Credentials): execution_project: Optional[str] = None quota_project: Optional[str] = None location: Optional[str] = None + api_endpoint: Optional[str] = None priority: Optional[Priority] = None maximum_bytes_billed: Optional[int] = None impersonate_service_account: Optional[str] = None diff --git a/dbt-bigquery/tests/unit/test_bigquery_adapter.py b/dbt-bigquery/tests/unit/test_bigquery_adapter.py index c3955cffc..2d97f8acd 100644 --- a/dbt-bigquery/tests/unit/test_bigquery_adapter.py +++ b/dbt-bigquery/tests/unit/test_bigquery_adapter.py @@ -84,6 +84,14 @@ def setUp(self): "priority": "batch", "maximum_bytes_billed": 0, }, + "api_endpoint": { + "type": "bigquery", + "method": "oauth", + "project": "dbt-unit-000000", + "schema": "dummy_schema", + "threads": 1, + "api_endpoint": "https://localhost:3001", + }, "impersonate": { "type": "bigquery", "method": "oauth", @@ -419,7 +427,7 @@ def test_cancel_open_connections_single(self): self.assertEqual(len(list(adapter.cancel_open_connections())), 1) @patch("dbt.adapters.bigquery.clients.ClientOptions") - @patch("dbt.adapters.bigquery.credentials.default") + @patch("dbt.adapters.bigquery.credentials._create_bigquery_defaults") @patch("dbt.adapters.bigquery.clients.BigQueryClient") def test_location_user_agent(self, MockClient, mock_auth_default, MockClientOptions): creds = MagicMock() @@ -439,6 +447,28 @@ def test_location_user_agent(self, MockClient, mock_auth_default, MockClientOpti client_options=mock_client_options, ) + @patch("dbt.adapters.bigquery.clients.ClientOptions") + @patch("dbt.adapters.bigquery.credentials._create_bigquery_defaults") + @patch("dbt.adapters.bigquery.clients.BigQueryClient") + def test_api_endpoint_settable(self, MockClient, mock_auth_default, MockClientOptions): + """Ensure api_endpoint is set on ClientOptions and passed to BigQueryClient.""" + + creds = MagicMock() + mock_auth_default.return_value = (creds, MagicMock()) + mock_client_options = MockClientOptions.return_value + + adapter = self.get_adapter("api_endpoint") + connection = adapter.acquire_connection("dummy") + MockClient.assert_not_called() + connection.handle + + MockClientOptions.assert_called_once() + kwargs = MockClientOptions.call_args.kwargs + assert kwargs.get("api_endpoint") == "https://localhost:3001" + + MockClient.assert_called_once() + assert MockClient.call_args.kwargs["client_options"] is mock_client_options + class HasUserAgent: PAT = re.compile(r"dbt-bigquery-\d+\.\d+\.\d+((a|b|rc)\d+)?")