From 8e1494dd1031df04298308ae7601eec3cc1e65a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20B=C3=BCcheler?= Date: Fri, 24 Jul 2026 14:37:02 +0200 Subject: [PATCH 1/2] fix(databricks): sanitize failed-task error text so it can be stored in jsonb MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Databricks run output is arbitrary external text and can contain NUL bytes or unpaired UTF-16 surrogates (e.g. when a task emits binary data). A failed task's error is placed in the DatabricksExecutionTrigger event and persisted by Airflow into the deferred task instance's next_kwargs (a jsonb column). Postgres jsonb rejects those characters, so the triggerer crashes on write (psycopg2 UntranslatableCharacter) and, because the invalid event is retried, enters a CrashLoopBackOff that stalls all deferrable tasks. Strip these characters from the extracted error in both the sync and async extract_failed_task_errors helpers before it is emitted in the event. Signed-off-by: Fabian Bücheler Co-authored-by: Cursor --- .../providers/databricks/utils/databricks.py | 28 ++++++- .../unit/databricks/utils/test_databricks.py | 76 +++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/utils/databricks.py b/providers/databricks/src/airflow/providers/databricks/utils/databricks.py index 05b6b17710ef8..112c304f3d476 100644 --- a/providers/databricks/src/airflow/providers/databricks/utils/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/utils/databricks.py @@ -17,9 +17,29 @@ # under the License. from __future__ import annotations +import re + from airflow.providers.common.compat.sdk import AirflowException, XComArg from airflow.providers.databricks.hooks.databricks import DatabricksHook, RunState +_JSONB_INVALID_CHARS = re.compile(r"[\x00\ud800-\udfff]") + + +def make_jsonb_safe(error: str | int) -> str | int: + """ + Strip characters that cannot be stored in a Postgres ``jsonb`` column from error text. + + Databricks run output is arbitrary external text and can contain NUL bytes or + unpaired UTF-16 surrogates (for example when a task emits binary data). These are + rejected by Postgres ``jsonb``, and since a failed task's error is persisted into + the deferred task instance's ``next_kwargs`` via the trigger event, an unsanitised + value crashes the triggerer when it writes the event. Non-string values are returned + unchanged. + """ + if isinstance(error, str): + return _JSONB_INVALID_CHARS.sub("", error) + return error + def normalise_json_content(content, json_path: str = "json") -> str | bool | list | dict | XComArg: """ @@ -75,7 +95,9 @@ def extract_failed_task_errors( error = run_output["error"] else: error = run_state.state_message - failed_tasks.append({"task_key": task_key, "run_id": task_run_id, "error": error}) + failed_tasks.append( + {"task_key": task_key, "run_id": task_run_id, "error": make_jsonb_safe(error)} + ) return failed_tasks @@ -101,7 +123,9 @@ async def extract_failed_task_errors_async( error = run_output["error"] else: error = run_state.state_message - failed_tasks.append({"task_key": task_key, "run_id": task_run_id, "error": error}) + failed_tasks.append( + {"task_key": task_key, "run_id": task_run_id, "error": make_jsonb_safe(error)} + ) return failed_tasks diff --git a/providers/databricks/tests/unit/databricks/utils/test_databricks.py b/providers/databricks/tests/unit/databricks/utils/test_databricks.py index 33716879713ed..61beca453da23 100644 --- a/providers/databricks/tests/unit/databricks/utils/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/utils/test_databricks.py @@ -27,6 +27,7 @@ from airflow.providers.databricks.utils.databricks import ( extract_failed_task_errors, extract_failed_task_errors_async, + make_jsonb_safe, normalise_json_content, validate_trigger_event, ) @@ -95,6 +96,18 @@ def test_validate_trigger_event_failure(self): with pytest.raises(AirflowException): validate_trigger_event(event) + def test_make_jsonb_safe_strips_nul_bytes(self): + assert make_jsonb_safe("PK\x03\x04binary\x00error") == "PK\x03\x04binaryerror" + + def test_make_jsonb_safe_strips_unpaired_surrogates(self): + assert make_jsonb_safe("bad\ud800surrogate\udfff") == "badsurrogate" + + def test_make_jsonb_safe_leaves_clean_text_unchanged(self): + assert make_jsonb_safe(ERROR_MESSAGE) == ERROR_MESSAGE + + def test_make_jsonb_safe_passes_through_non_strings(self): + assert make_jsonb_safe(123) == 123 + class TestExtractFailedTaskErrors: """Test cases for the extract_failed_task_errors utility function (synchronous version)""" @@ -178,6 +191,37 @@ def test_extract_failed_task_errors_single_failed_task_with_error_output(self, m assert result == expected hook.get_run_output.assert_called_once_with(TASK_RUN_ID_1) + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook") + def test_extract_failed_task_errors_sanitizes_error_output(self, mock_hook_class): + """Error text with jsonb-incompatible characters is sanitized so it can be persisted.""" + hook = mock_hook_class.return_value + hook.get_run_output = mock_dict({"error": "PK\x03\x04\x00broken\ud800"}) + + run_state = RunState("TERMINATED", "FAILED", "Job failed") + run_info = { + "run_id": RUN_ID, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "Job failed", + }, + "tasks": [ + { + "run_id": TASK_RUN_ID_1, + "task_key": TASK_KEY_1, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "Task failed", + }, + } + ], + } + + result = extract_failed_task_errors(hook, run_info, run_state) + + assert result == [{"task_key": TASK_KEY_1, "run_id": TASK_RUN_ID_1, "error": "PK\x03\x04broken"}] + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook") def test_extract_failed_task_errors_single_failed_task_without_error_output(self, mock_hook_class): """Test extracting errors from a single failed task without error in run output""" @@ -385,6 +429,38 @@ async def test_extract_failed_task_errors_async_single_failed_task(self, mock_ho assert result == expected hook.a_get_run_output.assert_called_once_with(TASK_RUN_ID_1) + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook") + @pytest.mark.asyncio + async def test_extract_failed_task_errors_async_sanitizes_error_output(self, mock_hook_class): + """Error text with jsonb-incompatible characters is sanitized so it can be persisted (async).""" + hook = mock_hook_class.return_value + hook.a_get_run_output = mock.AsyncMock(return_value={"error": "PK\x03\x04\x00broken\ud800"}) + + run_state = RunState("TERMINATED", "FAILED", "Job failed") + run_info = { + "run_id": RUN_ID, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "Job failed", + }, + "tasks": [ + { + "run_id": TASK_RUN_ID_1, + "task_key": TASK_KEY_1, + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "Task failed", + }, + } + ], + } + + result = await extract_failed_task_errors_async(hook, run_info, run_state) + + assert result == [{"task_key": TASK_KEY_1, "run_id": TASK_RUN_ID_1, "error": "PK\x03\x04broken"}] + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook") @pytest.mark.asyncio async def test_extract_failed_task_errors_async_multiple_failed_tasks(self, mock_hook_class): From 1b79983b513a20dfd4aa10991fe37faccdeb4f21 Mon Sep 17 00:00:00 2001 From: fabbuc-gyg <167207322+fabbuc-gyg@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:39:03 +0200 Subject: [PATCH 2/2] Update databricks.py change comment --- .../src/airflow/providers/databricks/utils/databricks.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/providers/databricks/src/airflow/providers/databricks/utils/databricks.py b/providers/databricks/src/airflow/providers/databricks/utils/databricks.py index 112c304f3d476..74938a3c66d0d 100644 --- a/providers/databricks/src/airflow/providers/databricks/utils/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/utils/databricks.py @@ -28,13 +28,6 @@ def make_jsonb_safe(error: str | int) -> str | int: """ Strip characters that cannot be stored in a Postgres ``jsonb`` column from error text. - - Databricks run output is arbitrary external text and can contain NUL bytes or - unpaired UTF-16 surrogates (for example when a task emits binary data). These are - rejected by Postgres ``jsonb``, and since a failed task's error is persisted into - the deferred task instance's ``next_kwargs`` via the trigger event, an unsanitised - value crashes the triggerer when it writes the event. Non-string values are returned - unchanged. """ if isinstance(error, str): return _JSONB_INVALID_CHARS.sub("", error)