[GitHub] [airflow] kukigai commented on a change in pull request #12126: Wait option for dagrun operator

[GitBox <gi...@apache.org>]

kukigai commented on a change in pull request #12126:
URL: https://github.com/apache/airflow/pull/12126#discussion_r518871327



##########
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 Good idea!  I updated the code and run tests.




----------------------------------------------------------------
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:
users@infra.apache.org