From f7aabed6d819031a67bef81ea871c0522fa7422e Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Thu, 1 May 2025 10:01:47 +0200 Subject: [PATCH 1/8] AddedSupport for api_endpoint config --- dbt-bigquery/src/dbt/adapters/bigquery/clients.py | 6 +++++- dbt-bigquery/src/dbt/adapters/bigquery/credentials.py | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py index 90c250dc5..9695fdb7a 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py @@ -56,12 +56,16 @@ def create_dataproc_batch_controller_client( @BQ_DEFAULT_RETRY def _create_bigquery_client(credentials: BigQueryCredentials) -> BigQueryClient: + kwargs_options = {} + if credentials.api_endpoint: + kwargs_options["api_endpoint"] = credentials.api_endpoint + return BigQueryClient( credentials.execution_project, 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, **kwargs_options), ) 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 From b8aca84334abba7b9811b0a69b4d74f2e6e80c2a Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Thu, 1 May 2025 08:19:23 +0000 Subject: [PATCH 2/8] ChanigeCommit --- .../.changes/unreleased/Features-20250501-081844.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 dbt-bigquery/.changes/unreleased/Features-20250501-081844.yaml 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 From 0e5599a4860e31fc1ca3355e44dca0d3b2897266 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Thu, 1 May 2025 08:53:27 +0000 Subject: [PATCH 3/8] Add tests. --- .../tests/unit/test_bigquery_adapter.py | 32 ++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/dbt-bigquery/tests/unit/test_bigquery_adapter.py b/dbt-bigquery/tests/unit/test_bigquery_adapter.py index c3955cffc..966740698 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 = get_adapter("api_endpoint") + 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+)?") From 68830198f8ecffd8d2cb5bfc02b7a20d8c972509 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Thu, 1 May 2025 08:56:25 +0000 Subject: [PATCH 4/8] Typo in Test --- dbt-bigquery/tests/unit/test_bigquery_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbt-bigquery/tests/unit/test_bigquery_adapter.py b/dbt-bigquery/tests/unit/test_bigquery_adapter.py index 966740698..d6e2da411 100644 --- a/dbt-bigquery/tests/unit/test_bigquery_adapter.py +++ b/dbt-bigquery/tests/unit/test_bigquery_adapter.py @@ -457,7 +457,7 @@ def test_api_endpoint_settable(self, MockClient, mock_auth_default, MockClientOp mock_auth_default.return_value = (creds, MagicMock()) mock_client_options = MockClientOptions.return_value - adapter = get_adapter("api_endpoint") + adapter = self.get_adapter("api_endpoint") adapter.acquire_connection("dummy") MockClient.assert_not_called() connection.handle From ececb3c944edbe09edc73e9f45735c16cd0ea474 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Thu, 1 May 2025 09:01:21 +0000 Subject: [PATCH 5/8] AnotherTypo --- dbt-bigquery/tests/unit/test_bigquery_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbt-bigquery/tests/unit/test_bigquery_adapter.py b/dbt-bigquery/tests/unit/test_bigquery_adapter.py index d6e2da411..2d97f8acd 100644 --- a/dbt-bigquery/tests/unit/test_bigquery_adapter.py +++ b/dbt-bigquery/tests/unit/test_bigquery_adapter.py @@ -458,7 +458,7 @@ def test_api_endpoint_settable(self, MockClient, mock_auth_default, MockClientOp mock_client_options = MockClientOptions.return_value adapter = self.get_adapter("api_endpoint") - adapter.acquire_connection("dummy") + connection = adapter.acquire_connection("dummy") MockClient.assert_not_called() connection.handle From 041f7c80ec8fc17edbf895ff905f088c512b1f1d Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Sun, 11 May 2025 14:53:08 +0200 Subject: [PATCH 6/8] Update dbt-bigquery/src/dbt/adapters/bigquery/clients.py Co-authored-by: Colin Rogers <111200756+colin-rogers-dbt@users.noreply.github.com> --- dbt-bigquery/src/dbt/adapters/bigquery/clients.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py index 9695fdb7a..6063a6c8b 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py @@ -65,7 +65,7 @@ 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, **kwargs_options), + client_options=ClientOptions(quota_project_id=credentials.quota_project, api_endpoint=credentials.api_endpoint), ) From 3e3703a20b0eaffb1156004ea037a61c00ea0001 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Mon, 12 May 2025 13:24:47 +0000 Subject: [PATCH 7/8] CleanUpDeadKwargs --- dbt-bigquery/src/dbt/adapters/bigquery/clients.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py index 6063a6c8b..db777329f 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py @@ -56,10 +56,6 @@ def create_dataproc_batch_controller_client( @BQ_DEFAULT_RETRY def _create_bigquery_client(credentials: BigQueryCredentials) -> BigQueryClient: - kwargs_options = {} - if credentials.api_endpoint: - kwargs_options["api_endpoint"] = credentials.api_endpoint - return BigQueryClient( credentials.execution_project, create_google_credentials(credentials), From 620196b9c2d282b7886ce40b2fc81fe75e90e078 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 12 May 2025 08:53:16 -0700 Subject: [PATCH 8/8] code formatting --- dbt-bigquery/src/dbt/adapters/bigquery/clients.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py index db777329f..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, api_endpoint=credentials.api_endpoint), + client_options=ClientOptions( + quota_project_id=credentials.quota_project, api_endpoint=credentials.api_endpoint + ), )