## Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you under the Apache License, Version 2.0 (the# "License"); you may not use this file except in compliance# with the License. You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing,# software distributed under the License is distributed on an# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY# KIND, either express or implied. See the License for the# specific language governing permissions and limitations# under the License.fromtypingimportIterable,Unionfromairflow.models.taskinstanceimportTaskInstancefromairflow.utilsimporttimezonefromairflow.utils.log.logging_mixinimportLoggingMixinfromairflow.utils.sessionimportcreate_session,provide_sessionfromairflow.utils.stateimportState# The key used by SkipMixin to store XCom data.
[docs]classSkipMixin(LoggingMixin):"""A Mixin to skip Tasks Instances"""
[docs]def_set_state_to_skipped(self,dag_run,execution_date,tasks,session):"""Used internally to set state of task instances to skipped from the same dag run."""task_ids=[d.task_idfordintasks]now=timezone.utcnow()ifdag_run:session.query(TaskInstance).filter(TaskInstance.dag_id==dag_run.dag_id,TaskInstance.execution_date==dag_run.execution_date,TaskInstance.task_id.in_(task_ids),).update({TaskInstance.state:State.SKIPPED,TaskInstance.start_date:now,TaskInstance.end_date:now,},synchronize_session=False,)else:ifexecution_dateisNone:raiseValueError("Execution date is None and no dag run")self.log.warning("No DAG RUN present this should not happen")# this is defensive against dag runs that are not completefortaskintasks:ti=TaskInstance(task,execution_date=execution_date)ti.state=State.SKIPPEDti.start_date=nowti.end_date=nowsession.merge(ti)
@provide_session
[docs]defskip(self,dag_run,execution_date,tasks,session=None,):""" Sets tasks instances to skipped from the same dag run. If this instance has a `task_id` attribute, store the list of skipped task IDs to XCom so that NotPreviouslySkippedDep knows these tasks should be skipped when they are cleared. :param dag_run: the DagRun for which to set the tasks to skipped :param execution_date: execution_date :param tasks: tasks to skip (not task_ids) :param session: db session to use """ifnottasks:returnself._set_state_to_skipped(dag_run,execution_date,tasks,session)session.commit()# SkipMixin may not necessarily have a task_id attribute. Only store to XCom if one is available.try:task_id=self.task_id# noqaexceptAttributeError:task_id=Noneiftask_idisnotNone:fromairflow.models.xcomimportXComXCom.set(key=XCOM_SKIPMIXIN_KEY,value={XCOM_SKIPMIXIN_SKIPPED:[d.task_idfordintasks]},task_id=task_id,dag_id=dag_run.dag_id,execution_date=dag_run.execution_date,session=session,
)
[docs]defskip_all_except(self,ti:TaskInstance,branch_task_ids:Union[str,Iterable[str]]):""" This method implements the logic for a branching operator; given a single task ID or list of task IDs to follow, this skips all other tasks immediately downstream of this operator. branch_task_ids is stored to XCom so that NotPreviouslySkippedDep knows skipped tasks or newly added tasks should be skipped when they are cleared. """self.log.info("Following branch %s",branch_task_ids)ifisinstance(branch_task_ids,str):branch_task_ids={branch_task_ids}branch_task_ids=set(branch_task_ids)dag_run=ti.get_dagrun()task=ti.taskdag=task.dagdownstream_tasks=task.downstream_listifdownstream_tasks:# For a branching workflow that looks like this, when "branch" does skip_all_except("task1"),# we intuitively expect both "task1" and "join" to execute even though strictly speaking,# "join" is also immediately downstream of "branch" and should have been skipped. Therefore,# we need a special case here for such empty branches: Check downstream tasks of branch_task_ids.# In case the task to skip is also downstream of branch_task_ids, we add it to branch_task_ids and# exclude it from skipping.## branch -----> join# \ ^# v /# task1#forbranch_task_idinlist(branch_task_ids):branch_task_ids.update(dag.get_task(branch_task_id).get_flat_relative_ids(upstream=False))skip_tasks=[tfortindownstream_tasksift.task_idnotinbranch_task_ids]follow_task_ids=[t.task_idfortindownstream_tasksift.task_idinbranch_task_ids]self.log.info("Skipping tasks %s",[t.task_idfortinskip_tasks])withcreate_session()assession:self._set_state_to_skipped(dag_run,ti.execution_date,skip_tasks,session=session)# For some reason, session.commit() needs to happen before xcom_push.# Otherwise the session is not committed.session.commit()ti.xcom_push(key=XCOM_SKIPMIXIN_KEY,value={XCOM_SKIPMIXIN_FOLLOWED:follow_task_ids})