From a96773a49fd9a7432d7ef86168bb8998fb64a090 Mon Sep 17 00:00:00 2001 From: jialuo Date: Fri, 4 Apr 2025 21:54:09 +0000 Subject: [PATCH 1/4] Add create_notebook_client in clients.py --- .../src/dbt/adapters/bigquery/clients.py | 14 ++++++++++++++ .../dbt/adapters/bigquery/python_submissions.py | 17 ++++++----------- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py index 722266240..bd03f4762 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py @@ -1,10 +1,12 @@ from google.api_core.client_info import ClientInfo from google.api_core.client_options import ClientOptions from google.auth.exceptions import DefaultCredentialsError +from google.cloud import aiplatform_v1 from google.cloud.bigquery import Client as BigQueryClient, DEFAULT_RETRY as BQ_DEFAULT_RETRY from google.cloud.dataproc_v1 import BatchControllerClient, JobControllerClient from google.cloud.storage import Client as StorageClient from google.cloud.storage.retry import DEFAULT_RETRY as GCS_DEFAULT_RETRY +from google.oauth2.credentials import Credentials as GoogleCredentials from dbt.adapters.events.logging import AdapterLogger @@ -67,3 +69,15 @@ def _create_bigquery_client(credentials: BigQueryCredentials) -> BigQueryClient: def _dataproc_endpoint(credentials: BigQueryCredentials) -> str: return f"{credentials.dataproc_region}-dataproc.googleapis.com:443" + + +def create_notebook_client( + credentials: GoogleCredentials, region: str | None +) -> aiplatform_v1.NotebookServiceClient: + api_endpoint = f"{region}-aiplatform.googleapis.com" + notebook_client = aiplatform_v1.NotebookServiceClient( + credentials=credentials, + client_options=ClientOptions(api_endpoint), + ) + + return notebook_client diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py b/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py index 7b42c2913..aa75f9c32 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py @@ -9,6 +9,7 @@ create_dataproc_batch_controller_client, create_dataproc_job_controller_client, create_gcs_client, + create_notebook_client ) from dbt.adapters.bigquery.credentials import ( BigQueryConnectionMethod, @@ -17,7 +18,6 @@ ) from dbt.adapters.bigquery.retry import RetryFactory from dbt.adapters.events.logging import AdapterLogger -from google.api_core.client_options import ClientOptions from google.auth.transport.requests import Request from google.cloud import aiplatform_v1 @@ -195,12 +195,7 @@ def __init__(self, parsed_model: Dict, credentials: BigQueryCredentials) -> None self._model_name = parsed_model["alias"] self._connection_method = credentials.method self._GoogleCredentials = create_google_credentials(credentials) - - # TODO(jialuo): Add a function in clients.py for it. - self._ai_platform_client = aiplatform_v1.NotebookServiceClient( - credentials=self._GoogleCredentials, - client_options=ClientOptions(api_endpoint=f"{self._region}-aiplatform.googleapis.com"), - ) + self._notebook_client = create_notebook_client(self._GoogleCredentials, self._region) self._notebook_template_id = parsed_model["config"].get("notebook_template_id") def _py_to_ipynb(self, compiled_code: str) -> str: @@ -220,7 +215,7 @@ def _get_notebook_template_id(self) -> str: parent=f"projects/{self._project}/locations/{self._region}", filter="notebookRuntimeType = ONE_CLICK", ) - page_result = self._ai_platform_client.list_notebook_runtime_templates(request=request) + page_result = self._notebook_client.list_notebook_runtime_templates(request=request) try: # Check if a default runtime template is available and applicable. @@ -269,7 +264,7 @@ def _create_notebook_template(self) -> str: notebook_runtime_template=template, ) - operation = self._ai_platform_client.create_notebook_runtime_template( + operation = self._notebook_client.create_notebook_runtime_template( request=create_request ) response = operation.result() @@ -408,7 +403,7 @@ def _submit_bigframes_job( ) try: - res = self._ai_platform_client.create_notebook_execution_job(request=request).result( + res = self._notebook_client.create_notebook_execution_job(request=request).result( timeout=self._polling_retry.timeout ) except TimeoutError as timeout_error: @@ -424,4 +419,4 @@ def _submit_bigframes_job( gcs_log_uri = f"{notebook_execution_job.gcs_output_uri}/{job_id}/{self._model_name}.py" self._process_gcs_log(gcs_log_uri) - return self._ai_platform_client.get_notebook_execution_job(name=res.name) + return self._notebook_client.get_notebook_execution_job(name=res.name) From e25c628b91a3f91f8628eff3823305969791f219 Mon Sep 17 00:00:00 2001 From: jialuo Date: Fri, 4 Apr 2025 21:56:57 +0000 Subject: [PATCH 2/4] changie --- dbt-bigquery/.changes/unreleased/Fixes-20250404-215640.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 dbt-bigquery/.changes/unreleased/Fixes-20250404-215640.yaml diff --git a/dbt-bigquery/.changes/unreleased/Fixes-20250404-215640.yaml b/dbt-bigquery/.changes/unreleased/Fixes-20250404-215640.yaml new file mode 100644 index 000000000..e03382c42 --- /dev/null +++ b/dbt-bigquery/.changes/unreleased/Fixes-20250404-215640.yaml @@ -0,0 +1,6 @@ +kind: Fixes +body: Add create_notebook_client +time: 2025-04-04T21:56:40.174432349Z +custom: + Author: jialuoo + Issue: "977" From cb0415306a3cdeca02216623be4086f75ca728d0 Mon Sep 17 00:00:00 2001 From: jialuo Date: Fri, 4 Apr 2025 22:06:49 +0000 Subject: [PATCH 3/4] quick fix --- .../src/dbt/adapters/bigquery/python_submissions.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py b/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py index aa75f9c32..1cba02ec4 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py @@ -9,7 +9,7 @@ create_dataproc_batch_controller_client, create_dataproc_job_controller_client, create_gcs_client, - create_notebook_client + create_notebook_client, ) from dbt.adapters.bigquery.credentials import ( BigQueryConnectionMethod, @@ -264,9 +264,7 @@ def _create_notebook_template(self) -> str: notebook_runtime_template=template, ) - operation = self._notebook_client.create_notebook_runtime_template( - request=create_request - ) + operation = self._notebook_client.create_notebook_runtime_template(request=create_request) response = operation.result() return self._extract_template_id(response.name) From b970dcbfd57dbea8dbd7a1de066af4a15376df58 Mon Sep 17 00:00:00 2001 From: jialuo Date: Fri, 4 Apr 2025 22:10:36 +0000 Subject: [PATCH 4/4] quick fix --- 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 bd03f4762..18714efcd 100644 --- a/dbt-bigquery/src/dbt/adapters/bigquery/clients.py +++ b/dbt-bigquery/src/dbt/adapters/bigquery/clients.py @@ -1,3 +1,5 @@ +from typing import Optional + from google.api_core.client_info import ClientInfo from google.api_core.client_options import ClientOptions from google.auth.exceptions import DefaultCredentialsError @@ -72,7 +74,7 @@ def _dataproc_endpoint(credentials: BigQueryCredentials) -> str: def create_notebook_client( - credentials: GoogleCredentials, region: str | None + credentials: GoogleCredentials, region: Optional[str] ) -> aiplatform_v1.NotebookServiceClient: api_endpoint = f"{region}-aiplatform.googleapis.com" notebook_client = aiplatform_v1.NotebookServiceClient(