kukigai commented on a change in pull request #12126:
URL: https://github.com/apache/airflow/pull/12126#discussion_r520807862
##########
File path: airflow/operators/dagrun_operator.py
##########
@@ -126,3 +148,58 @@ def execute(self, context: Dict):
dag.clear(start_date=self.execution_date,
end_date=self.execution_date)
else:
raise e
+
+ if self.wait_for_completion:
+ # wait for dag to complete
+ while True:
+ dttm = context['execution_date']
+
+ dttm_filter = dttm if isinstance(dttm, list) else [dttm]
+ serialized_dttm_filter = ','.join([datetime.isoformat() for
datetime in dttm_filter])
+
+ self.log.info(
+ 'Waiting for %s on %s to become one of %s... ',
+ self.trigger_dag_id,
+ serialized_dttm_filter,
+ self.allowed_states,
+ )
+
+ DR = DagRun
+
+ # get failed count
+ failed_count = (
+ session.query(func.count())
+ .filter(
+ DR.dag_id == self.trigger_dag_id,
+ DR.state.in_(self.failed_states), # pylint:
disable=no-member
+ DR.execution_date.in_(dttm_filter),
+ )
+ .scalar()
+ )
+
+ # if triggered dag run failed and that is not in
allowed_states,
+ # make this triggering dag fail.
+ if failed_count:
+ raise AirflowException(
+ f"{self.trigger_dag_id} failed with failed states
{self.failed_states}"
+ )
+
+ # get expected state count
+ count = (
+ session.query(func.count())
+ .filter(
+ DR.dag_id == self.trigger_dag_id,
+ DR.state.in_(self.allowed_states), # pylint:
disable=no-member
+ DR.execution_date.in_(dttm_filter),
+ )
+ .scalar()
+ )
+
+ session.commit()
Review comment:
@turbaszek
Can you please review the change?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
