Best way to trigger the external dag from master dag whose schedule interval = @once?

Viewed 2409

Having list of tasks which calls different dags from master dag.I'm using the TriggerDagrunoperator to accomplish this. But facing few issues.

  1. TriggerDagrunoperator doesn't wait for completion of external dag, it triggers next task. I want that to wait until completion and next task should trigger based on the status. Came across ExternalTaskSensor. It is making the process complicated. But the issue with externaltasksensor it only works with schedule. When i am using the externaltasksensor it is failing after timeout time has reached.
  2. Is there a way to trigger the different dags from a master dag sequentially not in parallel. So that it will trigger another dag only after the previously trigger dag has successfully completed all the tasks in the dag.
2 Answers

Run this task after the triggering your external dag:

import time
from airflow.models import DagRun
from airflow import AirflowException
from airflow.operators.python_operator import PythonOperator

def get_external_dag_status(dag_id,**kwargs):
    dag_id = dag_id 
    dag_runs = DagRun.find(dag_id=dag_id)
    for dag_run in dag_runs:
      #print("state = "+dag_run.state)
      res1 = dag_run.state
      #print(dag_run)
    return res1

def check_status(dag_id,**kwargs):
  st = get_external_dag_status(dag_id)
  while st != 'success':
    if st == 'failed':
      print(st)
      break
    time.sleep(300) #optional if need to check for every 5 minutes 
    st = get_external_dag_status(dag_id)
  if st == 'success':
    return st
  elif st == 'failed':
    raise ValueError('Dag Failed')

status_check = PythonOperator(task_id="dag_check",
                         python_callable=check_status,
                         op_kwargs={'dag_id':'your external dag id'},
                         dag=spark_dag
                         )

After triggering a Dag using TriggerDagrunoperator, you can consider calling a DagSensor, that will wait for the Dag completion, and only then trigger other days. Here how we implement our version (not perfect but did the job):

import logging

from airflow.plugins_manager import AirflowPlugin
from airflow.models import DagRun
from airflow.sensors.base_sensor_operator import BaseSensorOperator
from airflow.utils.db import provide_session
from airflow.utils.decorators import apply_defaults
from airflow.utils.state import State

logger = logging.getLogger('airflow.dag_sensor')


class DagSensor(BaseSensorOperator):
    """
    Sensor that check if a Dag is currently running.
    It proceeds only if the Dag is in not running.
    """

    template_fields = ['external_dag_id']
    ui_color = '#FFFFCC'

    @apply_defaults
    def __init__(self,
                 external_dag_id,
                 *args,
                 **kwargs):
        super(DagSensor, self).__init__(*args, **kwargs)
        self.external_dag_id = external_dag_id

    @provide_session
    def poke(self, context, session=None):
        dag_run = DagRun

        count = session.query(dag_run).filter(
            dag_run.dag_id == self.external_dag_id,
            dag_run._state.in_([State.RUNNING])
        ).count()
        session.commit()
        session.close()

        logger.info(f'Dag {self.external_dag_id} in running status: {count}')

        if count > 0:
            return False
        else:
            return True


class DagSensorPlugin(AirflowPlugin):
    name = 'dag_sensor_plugin'
    operators = [DagSensor]

Here how you can call it:

from airflow.operators import DagSensor
check_my_dag_completion = DagSensor(
    dag=dag,
    task_id='check_my_dag_completion',
    external_dag_id='my_dag',
    poke_interval=30,
    timeout=3600
)

This means that you can have something like this in your workflow:

call_dag_a >> check_dag_a >> call_dag_b >> check_dag_b
Related