ashb commented on code in PR #30478:
URL: https://github.com/apache/airflow/pull/30478#discussion_r1166547860
##########
airflow/models/dag.py:
##########
@@ -1915,9 +1914,118 @@ def set_task_instance_state(
only_failed=True,
session=session,
# Exclude the task itself from being cleared
- exclude_task_ids={task_id},
+ exclude_task_ids=frozenset({task_id}),
+ )
+
+ return altered
+
+ @provide_session
+ def set_task_group_state(
+ self,
+ *,
+ group_id: str,
+ execution_date: datetime | None = None,
+ run_id: str | None = None,
+ state: TaskInstanceState,
+ upstream: bool = False,
+ downstream: bool = False,
+ future: bool = False,
+ past: bool = False,
+ commit: bool = True,
+ session: Session = NEW_SESSION,
+ ) -> list[TaskInstance]:
+ """
+ Set the state of the TaskGroup to the given state, and clear its
downstream tasks that are
+ in failed or upstream_failed state.
+
+ :param group_id: The group_id of the TaskGroup
+ :param execution_date: Execution date of the TaskInstance
+ :param run_id: The run_id of the TaskInstance
+ :param state: State to set the TaskInstance to
+ :param upstream: Include all upstream tasks of the given task_id
+ :param downstream: Include all downstream tasks of the given task_id
+ :param future: Include all future TaskInstances of the given task_id
+ :param commit: Commit changes
+ :param past: Include all past TaskInstances of the given task_id
+ :param session: new session
+ """
+ from airflow.api.common.mark_tasks import set_state
+
+ if not exactly_one(execution_date, run_id):
+ raise ValueError("Exactly one of execution_date or run_id must be
provided")
+
+ tasks_to_set_state: list[BaseOperator | tuple[BaseOperator, int]] = []
+ task_ids: list[str] = []
+ locked_dag_run_ids: list[int] = []
+
+ if execution_date is None:
+ dag_run = (
+ session.query(DagRun).filter(DagRun.run_id == run_id,
DagRun.dag_id == self.dag_id).one()
+ ) # Raises an error if not found
+ resolve_execution_date = dag_run.execution_date
+ else:
+ resolve_execution_date = execution_date
+
+ end_date = resolve_execution_date if not future else None
+ start_date = resolve_execution_date if not past else None
+
+ task_group_dict = self.task_group.get_task_group_dict()
+ task_group = task_group_dict.get(group_id)
+ if task_group is None:
+ raise ValueError("TaskGroup {group_id} could not be found")
+ tasks_to_set_state = [task for task in task_group.iter_tasks() if
isinstance(task, BaseOperator)]
Review Comment:
What is this is-instance check guarding against? This will mean that mapped
tasks aren't cleared (as those don't inherit from BaseOperator). There is
`AbstractOperator` which is probably better to use than BaseOperator
I'd also think it would be good to not add the top level import -- all these
model files are intertwined enough, and adding imports to top levle can slow
down parsing/running tasks in unexpected ways
--
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]
