Follow-up of #47882 . In that PR TriggerDagRunOperator was ported to work with the Task SDK.
However, the deferral mode needs access to the DB, which we should port over.
and add similar logic to Airflow 2
|
# Kick off the deferral process |
|
if self._defer: |
|
self.defer( |
|
trigger=DagStateTrigger( |
|
dag_id=self.trigger_dag_id, |
|
states=self.allowed_states + self.failed_states, |
|
execution_dates=[dag_run.logical_date], |
|
run_ids=[run_id], |
|
poll_interval=self.poke_interval, |
|
), |
|
method_name="execute_complete", |
|
) |
at
|
@provide_session |
|
def execute_complete(self, context: Context, session: Session, event: tuple[str, dict[str, Any]]): |
|
# This logical_date is parsed from the return trigger event |
|
provided_logical_date = event[1]["execution_dates"][0] |
|
try: |
|
# Note: here execution fails on database isolation mode. Needs structural changes for AIP-72 |
|
dag_run = session.execute( |
|
select(DagRun).where( |
|
DagRun.dag_id == self.trigger_dag_id, DagRun.execution_date == provided_logical_date |
|
) |
|
).scalar_one() |
|
except NoResultFound: |
|
raise AirflowException( |
|
f"No DAG run found for DAG {self.trigger_dag_id} and logical date {self.logical_date}" |
|
) |
|
|
|
state = dag_run.state |
|
|
|
if state in self.failed_states: |
|
raise AirflowException(f"{self.trigger_dag_id} failed with failed state {state}") |
|
if state in self.allowed_states: |
|
self.log.info("%s finished with allowed state %s", self.trigger_dag_id, state) |
|
return |
|
|
|
raise AirflowException( |
|
f"{self.trigger_dag_id} return {state} which is not in {self.failed_states}" |
|
f" or {self.allowed_states}" |
|
) |
