## 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.from__future__importannotationsimportdatetimeimportosimportwarningsfromtypingimportTYPE_CHECKING,Any,Callable,Collection,Iterableimportattrfromairflow.configurationimportconffromairflow.exceptionsimportAirflowException,AirflowSkipException,RemovedInAirflow3Warningfromairflow.models.baseoperatorlinkimportBaseOperatorLinkfromairflow.models.dagimportDagModelfromairflow.models.dagbagimportDagBagfromairflow.models.taskinstanceimportTaskInstancefromairflow.operators.emptyimportEmptyOperatorfromairflow.sensors.baseimportBaseSensorOperatorfromairflow.triggers.external_taskimportWorkflowTriggerfromairflow.utils.fileimportcorrect_maybe_zippedfromairflow.utils.helpersimportbuild_airflow_url_with_queryfromairflow.utils.sensor_helperimport_get_count,_get_external_task_group_task_idsfromairflow.utils.sessionimportNEW_SESSION,provide_sessionfromairflow.utils.stateimportState,TaskInstanceStateifTYPE_CHECKING:fromsqlalchemy.ormimportSessionfromairflow.models.baseoperatorimportBaseOperatorfromairflow.models.taskinstancekeyimportTaskInstanceKeyfromairflow.utils.contextimportContext
[docs]classExternalDagLink(BaseOperatorLink):""" Operator link for ExternalTaskSensor and ExternalTaskMarker. It allows users to access DAG waited with ExternalTaskSensor or cleared by ExternalTaskMarker. """
[docs]classExternalTaskSensor(BaseSensorOperator):""" Waits for a different DAG, task group, or task to complete for a specific logical date. If both `external_task_group_id` and `external_task_id` are ``None`` (default), the sensor waits for the DAG. Values for `external_task_group_id` and `external_task_id` can't be set at the same time. By default, the ExternalTaskSensor will wait for the external task to succeed, at which point it will also succeed. However, by default it will *not* fail if the external task fails, but will continue to check the status until the sensor times out (thus giving you time to retry the external task without also having to clear the sensor). By default, the ExternalTaskSensor will not skip if the external task skips. To change this, simply set ``skipped_states=[TaskInstanceState.SKIPPED]``. Note that if you are monitoring multiple tasks, and one enters error state and the other enters a skipped state, then the external task will react to whichever one it sees first. If both happen together, then the failed state takes priority. It is possible to alter the default behavior by setting states which cause the sensor to fail, e.g. by setting ``allowed_states=[DagRunState.FAILED]`` and ``failed_states=[DagRunState.SUCCESS]`` you will flip the behaviour to get a sensor which goes green when the external task *fails* and immediately goes red if the external task *succeeds*! Note that ``soft_fail`` is respected when examining the failed_states. Thus if the external task enters a failed state and ``soft_fail == True`` the sensor will _skip_ rather than fail. As a result, setting ``soft_fail=True`` and ``failed_states=[DagRunState.SKIPPED]`` will result in the sensor skipping if the external task skips. However, this is a contrived example---consider using ``skipped_states`` if you would like this behaviour. Using ``skipped_states`` allows the sensor to skip if the target fails, but still enter failed state on timeout. Using ``soft_fail == True`` as above will cause the sensor to skip if the target fails, but also if it times out. :param external_dag_id: The dag_id that contains the task you want to wait for. (templated) :param external_task_id: The task_id that contains the task you want to wait for. (templated) :param external_task_ids: The list of task_ids that you want to wait for. (templated) If ``None`` (default value) the sensor waits for the DAG. Either external_task_id or external_task_ids can be passed to ExternalTaskSensor, but not both. :param external_task_group_id: The task_group_id that contains the task you want to wait for. (templated) :param allowed_states: Iterable of allowed states, default is ``['success']`` :param skipped_states: Iterable of states to make this task mark as skipped, default is ``None`` :param failed_states: Iterable of failed or dis-allowed states, default is ``None`` :param execution_delta: time difference with the previous execution to look at, the default is the same logical date as the current task or DAG. For yesterday, use [positive!] datetime.timedelta(days=1). Either execution_delta or execution_date_fn can be passed to ExternalTaskSensor, but not both. :param execution_date_fn: function that receives the current execution's logical date as the first positional argument and optionally any number of keyword arguments available in the context dictionary, and returns the desired logical dates to query. Either execution_delta or execution_date_fn can be passed to ExternalTaskSensor, but not both. :param check_existence: Set to `True` to check if the external task exists (when external_task_id is not None) or check if the DAG to wait for exists (when external_task_id is None), and immediately cease waiting if the external task or DAG does not exist (default value: False). :param poll_interval: polling period in seconds to check for the status :param deferrable: Run sensor in deferrable mode """
def__init__(self,*,external_dag_id:str,external_task_id:str|None=None,external_task_ids:Collection[str]|None=None,external_task_group_id:str|None=None,allowed_states:Iterable[str]|None=None,skipped_states:Iterable[str]|None=None,failed_states:Iterable[str]|None=None,execution_delta:datetime.timedelta|None=None,execution_date_fn:Callable|None=None,check_existence:bool=False,poll_interval:float=2.0,deferrable:bool=conf.getboolean("operators","default_deferrable",fallback=False),**kwargs,):super().__init__(**kwargs)self.allowed_states=list(allowed_states)ifallowed_stateselse[TaskInstanceState.SUCCESS.value]self.skipped_states=list(skipped_states)ifskipped_stateselse[]self.failed_states=list(failed_states)iffailed_stateselse[]total_states=set(self.allowed_states+self.skipped_states+self.failed_states)iflen(total_states)!=len(self.allowed_states)+len(self.skipped_states)+len(self.failed_states):raiseAirflowException("Duplicate values provided across allowed_states, skipped_states and failed_states.")# convert [] to Noneifnotexternal_task_ids:external_task_ids=None# can't set both single task id and a list of task idsifexternal_task_idisnotNoneandexternal_task_idsisnotNone:raiseValueError("Only one of `external_task_id` or `external_task_ids` may ""be provided to ExternalTaskSensor; ""use external_task_id or external_task_ids or external_task_group_id.")# since both not set, convert the single id to a 1-elt list - from here on, we only consider the listifexternal_task_idisnotNone:external_task_ids=[external_task_id]ifexternal_task_group_idisnotNoneandexternal_task_idsisnotNone:raiseValueError("Only one of `external_task_group_id` or `external_task_ids` may ""be provided to ExternalTaskSensor; ""use external_task_id or external_task_ids or external_task_group_id.")# check the requested states are all valid states for the target type, be it dag or taskifexternal_task_idsorexternal_task_group_id:ifnottotal_states<=set(State.task_states):raiseValueError("Valid values for `allowed_states`, `skipped_states` and `failed_states` ""when `external_task_id` or `external_task_ids` or `external_task_group_id` "f"is not `None`: {State.task_states}")elifnottotal_states<=set(State.dag_states):raiseValueError("Valid values for `allowed_states`, `skipped_states` and `failed_states` "f"when `external_task_id` and `external_task_group_id` is `None`: {State.dag_states}")ifexecution_deltaisnotNoneandexecution_date_fnisnotNone:raiseValueError("Only one of `execution_delta` or `execution_date_fn` may ""be provided to ExternalTaskSensor; not both.")self.execution_delta=execution_deltaself.execution_date_fn=execution_date_fnself.external_dag_id=external_dag_idself.external_task_id=external_task_idself.external_task_ids=external_task_idsself.external_task_group_id=external_task_group_idself.check_existence=check_existenceself._has_checked_existence=Falseself.deferrable=deferrableself.poll_interval=poll_intervaldef_get_dttm_filter(self,context):ifself.execution_delta:dttm=context["logical_date"]-self.execution_deltaelifself.execution_date_fn:dttm=self._handle_execution_date_fn(context=context)else:dttm=context["logical_date"]returndttmifisinstance(dttm,list)else[dttm]@provide_session
[docs]defpoke(self,context:Context,session:Session=NEW_SESSION)->bool:# delay check to poke rather than __init__ in case it was supplied as XComArgsifself.external_task_idsandlen(self.external_task_ids)>len(set(self.external_task_ids)):raiseValueError("Duplicate task_ids passed in external_task_ids parameter")dttm_filter=self._get_dttm_filter(context)serialized_dttm_filter=",".join(dt.isoformat()fordtindttm_filter)ifself.external_task_ids:self.log.info("Poking for tasks %s in dag %s on %s ... ",self.external_task_ids,self.external_dag_id,serialized_dttm_filter,)ifself.external_task_group_id:self.log.info("Poking for task_group '%s' in dag '%s' on %s ... ",self.external_task_group_id,self.external_dag_id,serialized_dttm_filter,)ifself.external_dag_idandnotself.external_task_group_idandnotself.external_task_ids:self.log.info("Poking for DAG '%s' on %s ... ",self.external_dag_id,serialized_dttm_filter,)# In poke mode this will check dag existence only onceifself.check_existenceandnotself._has_checked_existence:self._check_for_existence(session=session)count_failed=-1ifself.failed_states:count_failed=self.get_count(dttm_filter,session,self.failed_states)# Fail if anything in the list has failed.ifcount_failed>0:ifself.external_task_ids:ifself.soft_fail:raiseAirflowSkipException(f"Some of the external tasks {self.external_task_ids} "f"in DAG {self.external_dag_id} failed. Skipping due to soft_fail.")raiseAirflowException(f"Some of the external tasks {self.external_task_ids} "f"in DAG {self.external_dag_id} failed.")elifself.external_task_group_id:ifself.soft_fail:raiseAirflowSkipException(f"The external task_group '{self.external_task_group_id}' "f"in DAG '{self.external_dag_id}' failed. Skipping due to soft_fail.")raiseAirflowException(f"The external task_group '{self.external_task_group_id}' "f"in DAG '{self.external_dag_id}' failed.")else:ifself.soft_fail:raiseAirflowSkipException(f"The external DAG {self.external_dag_id} failed. Skipping due to soft_fail.")raiseAirflowException(f"The external DAG {self.external_dag_id} failed.")count_skipped=-1ifself.skipped_states:count_skipped=self.get_count(dttm_filter,session,self.skipped_states)# Skip if anything in the list has skipped. Note if we are checking multiple tasks and one skips# before another errors, we'll skip first.ifcount_skipped>0:ifself.external_task_ids:raiseAirflowSkipException(f"Some of the external tasks {self.external_task_ids} "f"in DAG {self.external_dag_id} reached a state in our states-to-skip-on list. Skipping.")elifself.external_task_group_id:raiseAirflowSkipException(f"The external task_group '{self.external_task_group_id}' "f"in DAG {self.external_dag_id} reached a state in our states-to-skip-on list. Skipping.")else:raiseAirflowSkipException(f"The external DAG {self.external_dag_id} reached a state in our states-to-skip-on list. ""Skipping.")# only go green if every single task has reached an allowed statecount_allowed=self.get_count(dttm_filter,session,self.allowed_states)returncount_allowed==len(dttm_filter)
[docs]defexecute(self,context:Context)->None:"""Run on the worker and defer using the triggers if deferrable is set to True."""ifnotself.deferrable:super().execute(context)else:self.defer(timeout=self.execution_timeout,trigger=WorkflowTrigger(external_dag_id=self.external_dag_id,external_task_group_id=self.external_task_group_id,external_task_ids=self.external_task_ids,execution_dates=self._get_dttm_filter(context),allowed_states=self.allowed_states,poke_interval=self.poll_interval,soft_fail=self.soft_fail,),method_name="execute_complete",)
[docs]defexecute_complete(self,context,event=None):"""Execute when the trigger fires - return immediately."""ifevent["status"]=="success":self.log.info("External tasks %s has executed successfully.",self.external_task_ids)elifevent["status"]=="skipped":raiseAirflowSkipException("External job has skipped skipping.")else:ifself.soft_fail:raiseAirflowSkipException("External job has failed skipping.")else:raiseAirflowException("Error occurred while trying to retrieve task status. Please, check the ""name of executed task and Dag.")
def_check_for_existence(self,session)->None:dag_to_wait=DagModel.get_current(self.external_dag_id,session)ifnotdag_to_wait:raiseAirflowException(f"The external DAG {self.external_dag_id} does not exist.")ifnotos.path.exists(correct_maybe_zipped(dag_to_wait.fileloc)):raiseAirflowException(f"The external DAG {self.external_dag_id} was deleted.")ifself.external_task_ids:refreshed_dag_info=DagBag(dag_to_wait.fileloc).get_dag(self.external_dag_id)forexternal_task_idinself.external_task_ids:ifnotrefreshed_dag_info.has_task(external_task_id):raiseAirflowException(f"The external task {external_task_id} in "f"DAG {self.external_dag_id} does not exist.")ifself.external_task_group_id:refreshed_dag_info=DagBag(dag_to_wait.fileloc).get_dag(self.external_dag_id)ifnotrefreshed_dag_info.has_task_group(self.external_task_group_id):raiseAirflowException(f"The external task group '{self.external_task_group_id}' in "f"DAG '{self.external_dag_id}' does not exist.")self._has_checked_existence=True
[docs]defget_count(self,dttm_filter,session,states)->int:""" Get the count of records against dttm filter and states. :param dttm_filter: date time filter for execution date :param session: airflow session object :param states: task or dag states :return: count of record against the filters """warnings.warn("This method is deprecated and will be removed in future.",DeprecationWarning,stacklevel=2)return_get_count(dttm_filter,self.external_task_ids,self.external_task_group_id,self.external_dag_id,states,session,)
[docs]defget_external_task_group_task_ids(self,session,dttm_filter):warnings.warn("This method is deprecated and will be removed in future.",DeprecationWarning,stacklevel=2)return_get_external_task_group_task_ids(dttm_filter,self.external_task_group_id,self.external_dag_id,session)
def_handle_execution_date_fn(self,context)->Any:""" Handle backward compatibility. This function is to handle backwards compatibility with how this operator was previously where it only passes the execution date, but also allow for the newer implementation to pass all context variables as keyword arguments, to allow for more sophisticated returns of dates to return. """fromairflow.utils.operator_helpersimportmake_kwargs_callable# Remove "logical_date" because it is already a mandatory positional argumentlogical_date=context["logical_date"]kwargs={k:vfork,vincontext.items()ifknotin{"execution_date","logical_date"}}# Add "context" in the kwargs for backward compatibility (because context used to be# an acceptable argument of execution_date_fn)kwargs["context"]=contextifTYPE_CHECKING:assertself.execution_date_fnisnotNonekwargs_callable=make_kwargs_callable(self.execution_date_fn)returnkwargs_callable(logical_date,**kwargs)
[docs]classExternalTaskMarker(EmptyOperator):""" Use this operator to indicate that a task on a different DAG depends on this task. When this task is cleared with "Recursive" selected, Airflow will clear the task on the other DAG and its downstream tasks recursively. Transitive dependencies are followed until the recursion_depth is reached. :param external_dag_id: The dag_id that contains the dependent task that needs to be cleared. :param external_task_id: The task_id of the dependent task that needs to be cleared. :param execution_date: The logical date of the dependent task execution that needs to be cleared. :param recursion_depth: The maximum level of transitive dependencies allowed. Default is 10. This is mostly used for preventing cyclic dependencies. It is fine to increase this number if necessary. However, too many levels of transitive dependencies will make it slower to clear tasks in the web UI. """
# The _serialized_fields are lazily loaded when get_serialized_fields() method is called__serialized_fields:frozenset[str]|None=Nonedef__init__(self,*,external_dag_id:str,external_task_id:str,execution_date:str|datetime.datetime|None="{{ logical_date.isoformat() }}",recursion_depth:int=10,**kwargs,):super().__init__(**kwargs)self.external_dag_id=external_dag_idself.external_task_id=external_task_idifisinstance(execution_date,datetime.datetime):self.execution_date=execution_date.isoformat()elifisinstance(execution_date,str):self.execution_date=execution_dateelse:raiseTypeError(f"Expected str or datetime.datetime type for execution_date. Got {type(execution_date)}")ifrecursion_depth<=0:raiseValueError("recursion_depth should be a positive integer")self.recursion_depth=recursion_depth@classmethod
[docs]defget_serialized_fields(cls):"""Serialize ExternalTaskMarker to contain exactly these fields + templated_fields ."""ifnotcls.__serialized_fields:cls.__serialized_fields=frozenset(super().get_serialized_fields()|{"recursion_depth"})returncls.__serialized_fields
@attr.s(auto_attribs=True)
[docs]classExternalTaskSensorLink(ExternalDagLink):""" This external link is deprecated. Please use :class:`airflow.sensors.external_task.ExternalDagLink`. """
[docs]def__attrs_post_init__(self):warnings.warn("This external link is deprecated. ""Please use :class:`airflow.sensors.external_task.ExternalDagLink`.",RemovedInAirflow3Warning,stacklevel=2,)