Skip to content

Commit

Permalink
Added support for custom api_endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
orlevii committed Nov 14, 2023
1 parent 0c5422c commit 7baa346
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -98,6 +99,7 @@ class BigQueryConnectionMethod(StrEnum):
SERVICE_ACCOUNT = "service-account"
SERVICE_ACCOUNT_JSON = "service-account-json"
OAUTH_SECRETS = "oauth-secrets"
ANONYMOUS = "anonymous"


@dataclass
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 7baa346

Please sign in to comment.