Repository: dbt-labs/dbt-adapters
Pull request #976: Raise an error when notebook execution failed
Author: login: jialuoo
Labels: name: cla:yes
description: The PR author has signed the CLA
Created: 2025-04-04T21:20:54Z | Merged: 2025-04-22T01:00:07Z

resolves #
[docs](https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose) dbt-labs/docs.getdbt.com/#

<!---
  Include the number of the issue addressed by this PR above if applicable.
  PRs for code changes without an associated issue *will not be merged*.
  See CONTRIBUTING.md for more information.

  Include the number of the docs issue that was opened for this PR. If
  this change has no user-facing implications, "N/A" suffices instead. New
  docs tickets can be created by clicking the link above or by going to
  https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose.
-->

### Problem

<!---
  Describe the problem this PR is solving. What is the application state
  before this PR is merged?
-->

### Solution

<!---
  Describe the way this PR solves the above problem. Add as much detail as you
  can to help reviewers understand your changes. Include any alternatives and
  tradeoffs you considered.
-->

### Checklist

- [ ] I have read [the contributing guide](https://github.com/dbt-labs/dbt-core/blob/main/CONTRIBUTING.md) and understand what's expected of me
- [ ] I have run this code in development and it appears to resolve the stated issue
- [ ] This PR includes tests, or tests are not required/relevant for this PR
- [ ] This PR has no interface changes (e.g. macros, cli, logs, json artifacts, config files, adapter interface, etc) or this PR has already received feedback and approval from Product or DX



Diff:
diff --git a/dbt-bigquery/.changes/unreleased/Fixes-20250404-212327.yaml b/dbt-bigquery/.changes/unreleased/Fixes-20250404-212327.yaml
new file mode 100644
index 000000000..01f126f3a
--- /dev/null
+++ b/dbt-bigquery/.changes/unreleased/Fixes-20250404-212327.yaml
@@ -0,0 +1,6 @@
+kind: Fixes
+body: Raise an error when notebook execution failed
+time: 2025-04-04T21:23:27.98150137Z
+custom:
+    Author: jialuoo
+    Issue: "976"
diff --git a/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py b/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
index 7051a9d8e..4218a4da1 100644
--- a/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
+++ b/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
@@ -17,10 +17,12 @@
 )
 from dbt.adapters.bigquery.retry import RetryFactory
 from dbt.adapters.events.logging import AdapterLogger
+from dbt_common.exceptions import DbtRuntimeError
 from google.api_core.client_options import ClientOptions
 
 from google.auth.transport.requests import Request
 from google.cloud import aiplatform_v1
+from google.cloud.aiplatform import gapic as aiplatform_gapic
 from google.cloud.dataproc_v1 import CreateBatchRequest, Job, RuntimeConfig
 from google.cloud.dataproc_v1.types.batches import Batch
 from google.protobuf.json_format import ParseDict
@@ -352,7 +354,14 @@ def _format_outputs(self, output_list: list) -> str:
                             formatted_output += f"    {inner_key}: {inner_value}\n"
                     else:
                         formatted_output += f"    {value}\n"
+
+                    # Extract only the final error message line. Discard the
+                    # traceback details as they contain raw ANSI escape codes.
+                    if isinstance(value, str) and value.strip().lower() == "error":
+                        return formatted_output
+
             else:
+                _logger.debug("Unexpected output format of the Colab notebook.")
                 formatted_output += f"{item}\n"
 
             # Add a newline between items.
@@ -410,6 +419,7 @@ def _submit_bigframes_job(
             res = self._ai_platform_client.create_notebook_execution_job(request=request).result(
                 timeout=self._polling_retry.timeout
             )
+            retrieved_job = self._ai_platform_client.get_notebook_execution_job(name=res.name)
         except TimeoutError as timeout_error:
             raise TimeoutError(
                 f"The dbt operation encountered a timeout: {timeout_error}\n"
@@ -417,10 +427,19 @@ def _submit_bigframes_job(
                 "console since it might still be actively running."
             )
         except Exception as e:
-            raise RuntimeError(f"An unexpected error occured while executing the notebook: {e}")
+            raise DbtRuntimeError(f"An unexpected error occured while executing the notebook: {e}")
 
         job_id = res.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:
+            raise DbtRuntimeError(
+                f"The colab notebook execution job '{retrieved_job.name}' failed."
+            )
+        elif retrieved_job.job_state != aiplatform_gapic.JobState.JOB_STATE_SUCCEEDED:
+            raise DbtRuntimeError(
+                f"The colab notebook execution job '{retrieved_job.name}' finished with unexpected state: {retrieved_job.job_state.name}"
+            )
+
         return self._ai_platform_client.get_notebook_execution_job(name=res.name)
diff --git a/dbt-bigquery/tests/functional/adapter/test_python_model.py b/dbt-bigquery/tests/functional/adapter/test_python_model.py
index fc8e59eac..f3e2119bf 100644
--- a/dbt-bigquery/tests/functional/adapter/test_python_model.py
+++ b/dbt-bigquery/tests/functional/adapter/test_python_model.py
@@ -306,6 +306,32 @@ def test_simple_bigframes_models(self, project):
         assert len(result) == 1
 
 
+models__bigframes_model_error = """
+def model(dbt, session):
+    dbt.config(
+        submission_method='bigframes',
+        materialized='table',
+    )
+    data = {"id": [1, 2, 3], "values": ['a', 'b', 'c']}
+    data += undefined_var
+    return bpd.DataFrame(data=data)
+"""
+
+
+@pytest.mark.flaky
+class TestBigframesModelsError:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "bigframes_model_error.py": models__bigframes_model_error,
+        }
+
+    def test_bigframes_models_error(self, project):
+        result, output = run_dbt_and_capture(["run"], expect_pass=False)
+        assert len(result) == 1
+        assert "name 'undefined_var' is not defined" in output
+
+
 models__bigframes_model_merge = """
 def model(dbt, session):
     dbt.config(


Reviews:
user: login: mikealfare
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T16:27:09Z
user: login: jialuoo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T19:20:38Z
user: login: jialuoo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T21:16:28Z
user: login: jialuoo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T23:56:41Z
user: login: mikealfare
body: This looks right, but I have one request. Can we add a test that shows that the error gets raise for a malformed model? I want to make sure that we fail if the user provides a malformed model.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-19T00:21:25Z
user: login: mikealfare
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-22T00:59:28Z

Review comments:
diff_hunk: @@ -352,7 +352,12 @@ def _format_outputs(self, output_list: list) -> str:
                             formatted_output += f"    {inner_key}: {inner_value}\n"
                     else:
                         formatted_output += f"    {value}\n"
+
+                    # Check errors from the output of notebook execution.
+                    if isinstance(value, str) and value.strip().lower() == "error":
path: dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
user: login: mikealfare
body: I'd rather not parse the output. Does the SDK not return the error? We should be catching errors, not parsing logs. If needed, can wee poll the job result like we do [here](https://github.com/dbt-labs/dbt-adapters/blob/5882b1df1f8f9ddcd0f4f5fcd09001b1948432e9/dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py#L109-L113)?
created_at: 2025-04-10T16:25:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -384,6 +389,8 @@ def _process_gcs_log(self, gcs_log_uri: str) -> None:
             # Improve the output format for better readability.
             formatted_output = self._format_outputs(outputs)
             _logger.info(f"Colab notebook runtime outputs from GCS: {formatted_output}")
+        except RuntimeError as e:
path: dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
user: login: mikealfare
body: This should be moved up prior to parsing the logs, and should raise a `DbtRuntimeError` from `e`.
created_at: 2025-04-10T16:26:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -352,7 +352,12 @@ def _format_outputs(self, output_list: list) -> str:
                             formatted_output += f"    {inner_key}: {inner_value}\n"
                     else:
                         formatted_output += f"    {value}\n"
+
+                    # Check errors from the output of notebook execution.
+                    if isinstance(value, str) and value.strip().lower() == "error":
path: dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
user: login: jialuoo
body: Yeah, the SDK doesn't return the error (it only has one state, <JobState.JOB_STATE_UNSPECIFIED: 0>, for all the cases either failure or success). So we have to parse the output to raise error.
created_at: 2025-04-10T19:20:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -384,6 +389,8 @@ def _process_gcs_log(self, gcs_log_uri: str) -> None:
             # Improve the output format for better readability.
             formatted_output = self._format_outputs(outputs)
             _logger.info(f"Colab notebook runtime outputs from GCS: {formatted_output}")
+        except RuntimeError as e:
path: dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
user: login: jialuoo
body: Yeah, DbtRuntimeError is better here. Fixed.
Not sure if I understand this `moved up prior to parsing the logs`. We need to parse the logs first to get the error message inside before raising the error.
created_at: 2025-04-10T21:16:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -352,7 +352,12 @@ def _format_outputs(self, output_list: list) -> str:
                             formatted_output += f"    {inner_key}: {inner_value}\n"
                     else:
                         formatted_output += f"    {value}\n"
+
+                    # Check errors from the output of notebook execution.
+                    if isinstance(value, str) and value.strip().lower() == "error":
path: dbt-bigquery/src/dbt/adapters/bigquery/python_submissions.py
user: login: jialuoo
body: Contacted the colab team, we could use GetNotebookExecutionJob RPC to do it (implemented).
created_at: 2025-04-18T23:56:41Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-04T21:21:40Z
body: Thank you for your pull request! We could not find a changelog entry for this change in the dbt-bigquery package. For details on how to document a change, see the [Contributing Guide](https://github.com/dbt-labs/dbt-adapters/blob/main/CONTRIBUTING.md).
author_association: CONTRIBUTOR
user: login: jialuoo
created_at: 2025-04-07T17:30:32Z
body: @mikealfare As the tests from this PR show: the setups for bigframes probably has some issue, which complains "location 'None' is set to an unknown value." Please help take a look
author_association: CONTRIBUTOR