PrakshiGoyal10 commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3949111289
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,39 @@ def get_run_tasks(self, run_id: int) -> list[dict[str,
Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a
terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's
metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act
on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED",
"MAXIMUM_CONCURRENT_RUNS_REACHED"}
+
+ # ``get_run_tasks`` returns one entry per attempt and is not ordered
by attempt, so a
+ # retried or already-repaired task appears several times under the
same ``task_key``. Keep
+ # only the latest attempt per key (sort by ``start_time``, same idiom
as
+ # ``DatabricksTaskBaseOperator._get_current_databricks_task``) before
judging its state, so
+ # the result never contains duplicate keys — Databricks rejects
duplicates in
+ # ``rerun_tasks`` — and a task whose latest attempt succeeded is not
reported as failed.
+ # Never-started sub-tasks omit ``start_time`` or send null; treat
those as 0.
+ sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda task:
task.get("start_time") or 0)
+ latest_by_key = {task["task_key"]: task for task in sorted_tasks}
+
+ failed_task_keys = []
+ for task_key, task in latest_by_key.items():
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
Review Comment:
The `dag_run` is used for two things: the clear step, and building the
redirect target from the run's own persisted `dag_id`/`run_id` rather than the
request path — that was to resolve the CodeQL open-redirect finding (a
same-site path from DB-sourced values, not request input). If I move clearing
to `dag.clear(run_id=...)`, the object would only be needed for that redirect.
Would you prefer I keep it DB-sourced for CodeQL, or drop the query and build
the redirect from the (validated) path params?
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @PrakshiGoyal10
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
