Repository: databricks/dbt-databricks
Pull request #985 (2025-04-05 11:19:42)

Fix: DBT Python Model Canceled Notebook Job Treated as Successful
<!-- Please review our pull request review process in CONTRIBUTING.md before your proceed. -->

Resolves #

https://github.com/databricks/dbt-databricks/issues/984

<!---
  Include the number of the issue addressed by this PR above if applicable.

  Example:
    resolves #1234

  Please review our pull request review process in CONTRIBUTING.md before your proceed.
-->

### Description
There is an issue where even after canceling a Notebook Job executed by the DBT Python model, it still appears successful. I have modified it to also check for CANCELLED in the terminal_states.

tested
```
hatch run unit
....
[gw1] [100%] PASSED tests/unit/test_adapter.py::TestDatabricksAdapter::test_describe_table_extended_may_limit 

===================== 469 passed, 6 skipped in 9.72s ======================
```


<!--- Describe the Pull Request here -->

### Checklist
- [x] I have run this code in development and it appears to resolve the stated issue
- [x] This PR includes tests, or tests are not required/relevant for this PR
- [x] I have updated the `CHANGELOG.md` and added information about my change to the "dbt-databricks next" section.


benc-db: Running functional tests now.  If the existing python functional tests pass, then I will give code one more look over and likely approve.
benc-db: Apologies, due to timing this got bumped out to 1.10.1; I do expect that release to come shortly, within next 2 weeks.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e82d3a0e..952bab38 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,9 @@
+## dbt-databricks 1.10.1 (TBD)
+
+### Fixes
+
+- Fix: DBT Python Model Canceled Notebook Job Treated as Successful ([985](https://github.com/databricks/dbt-databricks/pull/985))
+
 ## dbt-databricks 1.10.0 (Apr 08, 2025)
 
 ### Features
diff --git a/dbt/adapters/databricks/api_client.py b/dbt/adapters/databricks/api_client.py
index e2d34f1c..0f51f1a1 100644
--- a/dbt/adapters/databricks/api_client.py
+++ b/dbt/adapters/databricks/api_client.py
@@ -220,7 +220,7 @@ class PollableApi(DatabricksApi, ABC):
         params: dict,
         get_state_func: Callable[[Response], str],
         terminal_states: set[str],
-        expected_end_state: str,
+        expected_end_state: Optional[str],
         unexpected_end_state_func: Callable[[Response], None],
     ) -> Response:
         state = None
@@ -344,14 +344,20 @@ class JobRunsApi(PollableApi):
             params={"run_id": run_id},
             get_state_func=lambda response: response.json()["state"]["life_cycle_state"],
             terminal_states={"TERMINATED", "SKIPPED", "INTERNAL_ERROR"},
-            expected_end_state="TERMINATED",
+            expected_end_state=None,
             unexpected_end_state_func=self._get_exception,
         )
 
     def _get_exception(self, response: Response) -> None:
         response_json = response.json()
-        result_state = response_json["state"]["life_cycle_state"]
-        if result_state != "SUCCESS":
+        state = response_json["state"]
+        result_state = state.get("result_state")
+        life_cycle_state = state["life_cycle_state"]
+
+        if result_state is not None and result_state != "SUCCESS":
+            raise DbtRuntimeError(f"Python model run ended in result_state {result_state}")
+
+        if life_cycle_state != "TERMINATED":
             try:
                 task_id = response_json["tasks"][0]["run_id"]
                 # get end state to return to user
@@ -371,7 +377,7 @@ class JobRunsApi(PollableApi):
                 else:
                     state_message = response.json()["state"]["state_message"]
                     raise DbtRuntimeError(
-                        f"Python model run ended in state {result_state} "
+                        f"Python model run ended in state {life_cycle_state}"
                         f"with state_message\n{state_message}"
                     )
 


Test patch:
diff --git a/tests/unit/api_client/test_job_runs_api.py b/tests/unit/api_client/test_job_runs_api.py
index 9640136c..e7f9a978 100644
--- a/tests/unit/api_client/test_job_runs_api.py
+++ b/tests/unit/api_client/test_job_runs_api.py
@@ -84,6 +84,23 @@ class TestJobRunsApi(ApiTestBase):
         assert "Fancy exception" in str(exc.value)
         assert "trace" in str(exc.value)
 
+    @freezegun.freeze_time("2020-01-01")
+    @patch("dbt.adapters.databricks.api_client.time.sleep")
+    def test_poll_for_completion__cancelled(self, _, api, session):
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.return_value = {
+            "state": {
+                "life_cycle_state": "TERMINATED",
+                "result_state": "CANCELED",
+                "state_message": "cancelled by user",
+            }
+        }
+
+        with pytest.raises(DbtRuntimeError) as exc:
+            api.poll_for_completion("run_id")
+
+        assert "Python model run ended in result_state CANCELED" in str(exc.value)
+
     @freezegun.freeze_time("2020-01-01")
     @patch("dbt.adapters.databricks.api_client.time.sleep")
     def test_poll_for_completion__200(self, _, api, session, host):