How do i check if all my tasks in an airflow dag were successful?

Viewed 1054

i need to check if all the tasks of my dag were marked as successful so that in the last task of the dag it sends an email to me to notify if all were successful or if any failed.

here is a piece of code that i tried:

dag_runs = DagRun.find(dag_id=self.dagId)
for dag_run in dag_runs:
                
   if dag_run.state == 'success':
       body = f'\nHello , \nHere is the values for the pipeline {self.dagId}: \ncount of lines is {new_lines}, \nMax  date is {new_date}. \nRegards!'
   else: 
       body= f'\nHello  \nYour dag {self.dagId} has been Failed'   

email_text = """\
        Subject: %s    
        \nFrom: %s
        \nTo: %s
        
        \n%s
        """ % (subject, sent_from, self.to, body)

   
try:
   smtp_server = smtplib.SMTP_SSL('smtp.gmail.com', 465)
   smtp_server.ehlo()
   smtp_server.login(self.gmail_user, self.gmail_password)
   smtp_server.sendmail(sent_from, self.to, email_text)
   smtp_server.close()
   print ("Email sent successfully!")

except Exception as ex:
   print ("Something went wrong….",ex)

i'm unable to check if the dag state is success. so i want to check if the state of all tasks is success

thanks for the help and advice in advance.

2 Answers

We had a similar use-case where we want to identify if all the tasks were sucessful. In Airflow, if a task fails and if we have a trigger_rule one_failed, the DAG can run ends up being marked a successful as there was a recovery from failure.

Solution we implemented with single email to track all the task_instances:

from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance

def check_all_success(**context):
    dr: DagRun = context["dag_run"]
    ti: TaskInstance = context["ti"]
    
    # here we remove the task currently executing this logic
    ti_summary = set([task.state for task in dr.get_task_instances() if task.task_id != ti.task_id])

    # Remove success state
    ti_summary.remove('success')

    # If TI summary had any other state except success, there was an issue in the run
    if ti:
        # Send email: All tasks in DAG: {dr.dag_id} did not complete successfully
        pass
    else:
        # Send email: All tasks in DAG: {dr.dag_id} completed successfully
        pass

check_all_tasks = PythonOperator(
    task_id='check_all_tasks',
    python_callable=check_all_success,
    provide_context=True
)

By default, every task in Airflow should succeed for a next task to start running. So if your email-task is the last task in your DAG, that automatically means all previous tasks have succeeded.

Alternatively, you could configure on_success_callback and on_failure_callback on your DAG, which executes a given callable. This passes in arguments to determine whether the DAG run failed or succeeded:

def email(dagrun: DagRun, success: bool, reason: str, session: Session):
    # send email here...

success is a boolean value which indicates DAG Run success/failure.

Related