From 9622652d924cabf7818769bfe7be4f7457338e27 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Fri, 10 Nov 2023 16:40:06 +0100 Subject: [PATCH 1/8] Allow Bigquery Emulator settings to be set --- dbt/adapters/bigquery/connections.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index de84e4bf8..700c2f7ba 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 as 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, @@ -129,6 +129,7 @@ class BigQueryCredentials(Credentials): priority: Optional[Priority] = None maximum_bytes_billed: Optional[int] = None impersonate_service_account: Optional[str] = None + api_endpoint: Optional[str] = None job_retry_deadline_seconds: Optional[int] = None job_retries: Optional[int] = 1 @@ -381,6 +382,11 @@ def get_bigquery_client(cls, profile_credentials): creds = cls.get_credentials(profile_credentials) execution_project = profile_credentials.execution_project location = getattr(profile_credentials, "location", None) + options: Optional[client_options.ClientOptions] = None + if profile_credentials.api_endpoint: + options = client_options.ClientOptions( + api_endpoint=profile_credentials.api_endpoint, + ) info = client_info.ClientInfo(user_agent=f"dbt-{dbt_version}") return google.cloud.bigquery.Client( @@ -388,6 +394,7 @@ def get_bigquery_client(cls, profile_credentials): creds, location=location, client_info=info, + client_options=options, ) @classmethod From c13d264255e659b6f83f0a655f23448280989dfd Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Sat, 11 Nov 2023 09:08:16 +0100 Subject: [PATCH 2/8] Switch Connection ordering on api_enpoint --- dbt/adapters/bigquery/connections.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index 700c2f7ba..2a3f88ec9 100644 --- a/dbt/adapters/bigquery/connections.py +++ b/dbt/adapters/bigquery/connections.py @@ -126,10 +126,10 @@ class BigQueryCredentials(Credentials): schema: Optional[str] = None # type: ignore execution_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 - api_endpoint: Optional[str] = None job_retry_deadline_seconds: Optional[int] = None job_retries: Optional[int] = 1 From b2aa0dfe0f6f4dc302ff39a66062c1ff712a55e3 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Tue, 27 Feb 2024 17:11:08 +0100 Subject: [PATCH 3/8] Basic test to lock in that "api_endpoint" is a credentials field that should be passed on to BQ --- tests/unit/test_bigquery_adapter.py | 43 +++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index 34abd0caf..07e6d74d6 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -71,6 +71,17 @@ def setUp(self): "priority": "batch", "maximum_bytes_billed": 0, }, + "api_endpoint": { + "type": "bigquery", + "method": "oauth", + "project": "dbt-unit-000000", + "schema": "dummy_schema", + "threads": 1, + "location": "Luna Station", + "priority": "batch", + "maximum_bytes_billed": 0, + "api_endpoint: "https://localhost:3001", + }, "impersonate": { "type": "bigquery", "method": "oauth", @@ -406,6 +417,38 @@ def test_location_user_agent(self, mock_bq, mock_auth_default): client_info=HasUserAgent(), ) + @patch("dbt.adapters.bigquery.impl.google.auth.default") + @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery") + def test_api_endpoint_settable(self, mock_bq, mock_auth_default): + """Ensure that a user can pass an endpoint to the connector such as is needed for an emulator.""" + + creds = MagicMock() + mock_auth_default.return_value = (creds, MagicMock()) + adapter = self.get_adapter("api_endpoint") + + connection = adapter.acquire_connection("dummy") + mock_client = mock_bq.Client + + mock_client.assert_not_called() + connection.handle + mock_client.assert_called_once_with( + "dbt-unit-000000", + creds, + location="Luna Station", + client_info=HasUserAgent(), + client_options=_CheckApiEndpointSet("https://localhost:3001") + ) + + +class _CheckApiEndpointSet: + value: str + + def __init__(self, value: str) -> None: + self.value = value + + def __eq__(self, other) -> bool: + return getattr(other, "api_endpoint") == self.value + class HasUserAgent: PAT = re.compile(r"dbt-bigquery-\d+\.\d+\.\d+((a|b|rc)\d+)?") From 03b6a6d52091167f2a9ed811cbbb3a9b4c65aceb Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Tue, 27 Feb 2024 17:13:13 +0100 Subject: [PATCH 4/8] Example ApiEndpoint typo --- tests/unit/test_bigquery_adapter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index 07e6d74d6..aff04581c 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -80,7 +80,7 @@ def setUp(self): "location": "Luna Station", "priority": "batch", "maximum_bytes_billed": 0, - "api_endpoint: "https://localhost:3001", + "api_endpoint": "https://localhost:3001", }, "impersonate": { "type": "bigquery", @@ -420,7 +420,7 @@ def test_location_user_agent(self, mock_bq, mock_auth_default): @patch("dbt.adapters.bigquery.impl.google.auth.default") @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery") def test_api_endpoint_settable(self, mock_bq, mock_auth_default): - """Ensure that a user can pass an endpoint to the connector such as is needed for an emulator.""" + """Ensure that a user can pass api_endpoint to the connector eg. for emulator.""" creds = MagicMock() mock_auth_default.return_value = (creds, MagicMock()) From 9478ae9f8d4bf6a1be5487ed23d5c8e4330405bb Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Mon, 18 Mar 2024 22:16:17 +0100 Subject: [PATCH 5/8] Linting/Black bump --- tests/unit/test_bigquery_adapter.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index aff04581c..31347ca83 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -436,16 +436,16 @@ def test_api_endpoint_settable(self, mock_bq, mock_auth_default): creds, location="Luna Station", client_info=HasUserAgent(), - client_options=_CheckApiEndpointSet("https://localhost:3001") + client_options=_CheckApiEndpointSet("https://localhost:3001"), ) class _CheckApiEndpointSet: value: str - + def __init__(self, value: str) -> None: self.value = value - + def __eq__(self, other) -> bool: return getattr(other, "api_endpoint") == self.value From a9ee09853f2215b4b36480ec25ce06beeb4c3976 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Mon, 18 Mar 2024 22:18:52 +0100 Subject: [PATCH 6/8] Correct call signature --- tests/unit/test_bigquery_adapter.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index 31347ca83..0f42574ec 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -415,6 +415,7 @@ def test_location_user_agent(self, mock_bq, mock_auth_default): creds, location="Luna Station", client_info=HasUserAgent(), + client_options=None, ) @patch("dbt.adapters.bigquery.impl.google.auth.default") From 28009c91380c98d0cecbe598f0f3a0c22d4bcdf6 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Sun, 24 Mar 2024 18:47:36 +0000 Subject: [PATCH 7/8] Break out client mock args for clarity --- tests/unit/test_bigquery_adapter.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index 0f42574ec..3aea61115 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -410,13 +410,12 @@ def test_location_user_agent(self, mock_bq, mock_auth_default): mock_client.assert_not_called() connection.handle - mock_client.assert_called_once_with( - "dbt-unit-000000", - creds, - location="Luna Station", - client_info=HasUserAgent(), - client_options=None, - ) + # Get the arguments passed to mock_client.assert_called_once_with() + call_args = mock_client.call_args + args, kwargs = call_args + assert args == ("dbt-unit-000000", creds) + assert kwargs["location"] == "Luna Station" + assert HasUserAgent() == kwargs["client_info"] @patch("dbt.adapters.bigquery.impl.google.auth.default") @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery") From 5d51658929554979e1efb72c063362343a59a264 Mon Sep 17 00:00:00 2001 From: Michael O Toole Date: Mon, 25 Mar 2024 10:44:55 +0000 Subject: [PATCH 8/8] Clear caching between tests --- tests/unit/test_bigquery_adapter.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_bigquery_adapter.py b/tests/unit/test_bigquery_adapter.py index 3aea61115..9f906f48a 100644 --- a/tests/unit/test_bigquery_adapter.py +++ b/tests/unit/test_bigquery_adapter.py @@ -17,7 +17,7 @@ from dbt.adapters.bigquery.relation_configs import PartitionConfig from dbt.adapters.bigquery import BigQueryAdapter, BigQueryRelation from google.cloud.bigquery.table import Table -from dbt.adapters.bigquery.connections import _sanitize_label, _VALIDATE_LABEL_LENGTH_LIMIT +from dbt.adapters.bigquery.connections import _sanitize_label, _VALIDATE_LABEL_LENGTH_LIMIT, get_bigquery_defaults from dbt_common.clients import agate_helper import dbt_common.exceptions from dbt.context.manifest import generate_query_header_context @@ -401,6 +401,7 @@ def test_cancel_open_connections_single(self): @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): + get_bigquery_defaults.cache_clear() creds = MagicMock() mock_auth_default.return_value = (creds, MagicMock()) adapter = self.get_adapter("loc") @@ -410,18 +411,20 @@ def test_location_user_agent(self, mock_bq, mock_auth_default): mock_client.assert_not_called() connection.handle - # Get the arguments passed to mock_client.assert_called_once_with() - call_args = mock_client.call_args - args, kwargs = call_args - assert args == ("dbt-unit-000000", creds) - assert kwargs["location"] == "Luna Station" - assert HasUserAgent() == kwargs["client_info"] + mock_client.assert_called_once_with( + "dbt-unit-000000", + creds, + location="Luna Station", + client_info=HasUserAgent(), + client_options=None, + ) @patch("dbt.adapters.bigquery.impl.google.auth.default") @patch("dbt.adapters.bigquery.impl.google.cloud.bigquery") def test_api_endpoint_settable(self, mock_bq, mock_auth_default): """Ensure that a user can pass api_endpoint to the connector eg. for emulator.""" + get_bigquery_defaults.cache_clear() creds = MagicMock() mock_auth_default.return_value = (creds, MagicMock()) adapter = self.get_adapter("api_endpoint")