tquazi commented on code in PR #30478:
URL: https://github.com/apache/airflow/pull/30478#discussion_r1167622973
##########
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:
task_group.iter_tasks returns an iterator of AbstractOperator objects.
I am using is-instance to cast AbstractOperator to BaseOperator.
The set_state method expects List<BaseOperator> as argument and the
TypeChecking is failing this piece of code.
https://github.com/apache/airflow/pull/30478/files#diff-62c8e300ee91e0d59f81e0ea5d30834f04db71ae74f2e155a10b51056b00b59bR1991-R2002
We have similar use case in views.py But looks like TypeChecking is not
enabled here.
https://github.com/apache/airflow/blob/main/airflow/www/views.py#L2183
Any suggestions to resolve this?
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
