Track the notebook job status to avoid LRO (#1175)

Co-authored-by: Colin Rogers <[email protected]>

Author: jialuoo

dbt-bigquery/.changes/unreleased/Fixes-20250626-171343.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: Track notebook job status to avoid LRO in bigframes mode
+time: 2025-06-26T17:13:43.757335863Z
+custom:
+    Author: jialuoo
+    Issue: "1175"

dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py

@@ -1,6 +1,7 @@
 import inspect
 import json
 import re
+import time
 from typing import Any, Dict, Optional, Union
 import uuid
 
@@ -21,6 +22,7 @@
 from dbt.adapters.events.logging import AdapterLogger
 from dbt_common.exceptions import DbtRuntimeError
 
+from google.api_core.operation import Operation
 from google.auth.transport.requests import Request
 from google.cloud import aiplatform_v1
 from google.cloud.aiplatform import gapic as aiplatform_gapic
@@ -42,6 +44,9 @@
 # This differs from other Python models because the typical 5-minute timeout
 # is insufficient for BigFrames processing.
 _DEFAULT_BIGFRAMES_TIMEOUT = 60 * 60
+# Time interval in seconds between successive polling attempts to check the
+# notebook job's status in BigFrames mode.
+_COLAB_POLL_INTERVAL = 30
 
 
 class _BigQueryPythonHelper(PythonJobHelper):
@@ -415,6 +420,39 @@ def submit(self, compiled_code: str) -> None:
 
         self._submit_bigframes_job(notebook_template_id)
 
+    def _track_notebook_job_status(self, job_name: str) -> aiplatform_v1.NotebookExecutionJob:
+        """Tracks the notebook job until it completes or times out."""
+        max_wait_time = self._polling_retry.timeout
+        elapsed = 0
+
+        # Please see all the JobState from
+        # https://cloud.google.com/php/docs/reference/cloud-ai-platform/latest/V1.JobState.
+        terminal_states = {
+            aiplatform_gapic.JobState.JOB_STATE_SUCCEEDED,
+            aiplatform_gapic.JobState.JOB_STATE_PARTIALLY_SUCCEEDED,
+            aiplatform_gapic.JobState.JOB_STATE_FAILED,
+            aiplatform_gapic.JobState.JOB_STATE_CANCELLED,
+            aiplatform_gapic.JobState.JOB_STATE_EXPIRED,
+        }
+
+        while True:
+            retrieved_job = self._notebook_client.get_notebook_execution_job(name=job_name)
+            job_state = retrieved_job.job_state
+
+            if job_state in terminal_states:
+                return retrieved_job
+
+            if elapsed >= max_wait_time:
+                raise TimeoutError(
+                    "Operation did not complete within the designated timeout "
+                    f"of {max_wait_time} seconds. Please cancel the related "
+                    "notebook job manually via the GCP console since it might "
+                    "still be actively running."
+                )
+
+            time.sleep(_COLAB_POLL_INTERVAL)
+            elapsed += _COLAB_POLL_INTERVAL
+
     def _submit_bigframes_job(
         self, notebook_template_id: str
     ) -> aiplatform_v1.NotebookExecutionJob:
@@ -427,33 +465,36 @@ def _submit_bigframes_job(
         )
 
         try:
-            res = self._notebook_client.create_notebook_execution_job(request=request).result(
-                timeout=self._polling_retry.timeout
+            operation: Operation = self._notebook_client.create_notebook_execution_job(
+                request=request
             )
-            retrieved_job = self._notebook_client.get_notebook_execution_job(name=res.name)
+            lro_name = operation.operation.name
+            job_name = lro_name.split("/operations/")[0]
+            retrieved_job = self._track_notebook_job_status(job_name)
         except TimeoutError as timeout_error:
-            raise TimeoutError(
-                f"The dbt operation encountered a timeout: {timeout_error}\n"
-                "Please cancel the related notebook job manually via the GCP "
-                "console since it might still be actively running."
-            )
+            raise TimeoutError(f"The dbt operation encountered a timeout: {timeout_error}")
         except Exception as e:
             raise DbtRuntimeError(f"An unexpected error occured while executing the notebook: {e}")
 
-        job_id = res.name.split("/")[-1]
+        job_id = job_name.split("/")[-1]
         gcs_log_uri = f"{notebook_execution_job.gcs_output_uri}/{job_id}/{self._model_name}.py"
         self._process_gcs_log(gcs_log_uri)
 
-        if retrieved_job.job_state == aiplatform_gapic.JobState.JOB_STATE_FAILED:
+        if retrieved_job.job_state == aiplatform_gapic.JobState.JOB_STATE_SUCCEEDED:
+            _logger.info(
+                f"Colab notebook execution job '{retrieved_job.name}' finished successfully."
+            )
+        elif retrieved_job.job_state == aiplatform_gapic.JobState.JOB_STATE_FAILED:
             raise DbtRuntimeError(
                 f"The colab notebook execution job '{retrieved_job.name}' failed."
             )
-        elif retrieved_job.job_state != aiplatform_gapic.JobState.JOB_STATE_SUCCEEDED:
+        else:
             raise DbtRuntimeError(
-                f"The colab notebook execution job '{retrieved_job.name}' finished with unexpected state: {retrieved_job.job_state.name}"
+                f"The colab notebook execution job '{retrieved_job.name}' "
+                f"finished with unexpected state: {retrieved_job.job_state.name}"
             )
 
-        return self._notebook_client.get_notebook_execution_job(name=res.name)
+        return self._notebook_client.get_notebook_execution_job(name=job_name)
 
 
 def _install_packages(packages: list[str]) -> None: