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

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

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



##########
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:
       Should we do:
   ```py
   dag_run = trigger_dagrun(...) # In L126
   ...
   if self._wait_for_completion:
       state = None    
       while True:
             dag_run.refresh_from_db()
             state = dag_run.state
             if state in self.failed_states:
                   raise AirflowException(...)
             elif state in self.allowed_states:
                   self.log.info(...)
                   return
             sleep(...)
   ```
   WDYT?




----------------------------------------------------------------
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