From 419d8c62b1491a85a9c80c29e6d7173ffd3ca1ce Mon Sep 17 00:00:00 2001 From: raghvendra-singh1 Date: Wed, 11 Dec 2024 14:02:09 +0530 Subject: [PATCH] [FIX] Fixed databricks repair run deferrable (#44213) * [FIX] Fixed databricks repair run deferrable * Removed trailing spaces as per pre-commit hook * [Fix] Modified code to handle invalid trigger event --- .../providers/databricks/operators/databricks.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/providers/src/airflow/providers/databricks/operators/databricks.py b/providers/src/airflow/providers/databricks/operators/databricks.py index 2f9a460703396..1b8d45fa479a3 100644 --- a/providers/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/src/airflow/providers/databricks/operators/databricks.py @@ -930,13 +930,15 @@ def execute(self, context: Context): def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None: if event: + if event.get("run_state"): + run_state = RunState.from_json(event["run_state"]) + if event.get("repair_run"): + event["repair_run"] = event["repair_run"] and ( + not self.databricks_repair_reason_new_settings + or is_repair_reason_match_exist(self, run_state) + ) _handle_deferrable_databricks_operator_completion(event, self.log) - run_state = RunState.from_json(event["run_state"]) - should_repair = event["repair_run"] and ( - not self.databricks_repair_reason_new_settings - or is_repair_reason_match_exist(self, run_state) - ) - if should_repair: + if event.get("repair_run"): self.repair_run = False self.run_id = event["run_id"] job_id = self._hook.get_job_id(self.run_id)