From 7baa346b65f1f8144a3cfa8e91c05d5f4a9be6e0 Mon Sep 17 00:00:00 2001 From: Or Levi Date: Tue, 14 Nov 2023 16:49:55 +0000 Subject: [PATCH] Added support for custom api_endpoint --- dbt/adapters/bigquery/connections.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/dbt/adapters/bigquery/connections.py b/dbt/adapters/bigquery/connections.py index de84e4bf8..06f0af4cd 100644 --- a/dbt/adapters/bigquery/connections.py +++ b/dbt/adapters/bigquery/connections.py @@ -17,8 +17,9 @@ 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.auth.credentials import AnonymousCredentials from google.oauth2 import ( credentials as GoogleCredentials, service_account as GoogleServiceAccountCredentials, @@ -98,6 +99,7 @@ class BigQueryConnectionMethod(StrEnum): SERVICE_ACCOUNT = "service-account" SERVICE_ACCOUNT_JSON = "service-account-json" OAUTH_SECRETS = "oauth-secrets" + ANONYMOUS = "anonymous" @dataclass @@ -355,6 +357,8 @@ def get_google_credentials(cls, profile_credentials) -> GoogleCredentials: token_uri=profile_credentials.token_uri, scopes=profile_credentials.scopes, ) + elif method == BigQueryConnectionMethod.ANONYMOUS: + return AnonymousCredentials() error = 'Invalid `method` in profile: "{}"'.format(method) raise FailedToConnectError(error) @@ -383,11 +387,11 @@ def get_bigquery_client(cls, profile_credentials): location = getattr(profile_credentials, "location", None) info = client_info.ClientInfo(user_agent=f"dbt-{dbt_version}") + + api_endpoint = getattr(profile_credentials, "api_endpoint", None) + options = client_options.ClientOptions(api_endpoint=api_endpoint) return google.cloud.bigquery.Client( - execution_project, - creds, - location=location, - client_info=info, + execution_project, creds, location=location, client_info=info, client_options=options ) @classmethod