10

I have implemented email alerts on success and failure using on_success_callback and on_failure_callback.

According to Airflow documentation,

a context dictionary is passed as a single parameter to this function.

How can I pass another parameter to these callback methods?

Here is my code

from airflow.utils.email import send_email_smtp

def task_success_alert(context):
    subject = "[Airflow] DAG {0} - Task {1}: Success".format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1]
        )
    html_content = """
    DAG: {0}<br>
    Task: {1}<br>
    Succeeded on: {2}
    """.format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1], 
        datetime.now()
        )
    send_email_smtp(dag_vars["dev_mailing_list"], subject, html_content)

def task_failure_alert(context):
    subject = "[Airflow] DAG {0} - Task {1}: Failed".format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1]
        )
    html_content = """
    DAG: {0}<br>
    Task: {1}<br>
    Failed on: {2}
    """.format(
        context['task_instance_key_str'].split('__')[0], 
        context['task_instance_key_str'].split('__')[1], 
        datetime.now()
        )
    send_email_smtp(dag_vars["dev_mailing_list"], subject, html_content)

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2019, 6, 13),
    'on_success_callback': task_success_alert,
    'on_failure_callback': task_failure_alert
}

I intend to move the callbacks to another package and pass the email address as parameter.

Blessy
  • 480
  • 2
  • 7
  • 20
  • When I use CONTEXT it my tasks finishes but dag keeps running.It never ends. any suggestions. I am using def task_alert(context): dag_id = context['dag'].dag_id \n task_id = context['task_instance']. task_id I am calling on_failure_callback=task_alert – Akshay Lande Jun 21 '20 at 06:37
  • Use the xcom_push and xcom_pull – jetsun Oct 15 '20 at 05:40

3 Answers3

10

You could define a function inside your dag that calls the function from your package. And while calling that function, pass email as an argument. You can refine it further at your DAG level to pass only information required for the emails.

from package import outer_task_success_callback
email = 'xyz@example.com'

def task_success_alert(context):
    dag_id = context['dag'].dag_id
    task_id = context['task_instance']. task_id
    outer_task_success_callback(dag_id, task_id, email)
    
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2019, 6, 13),
    'on_success_callback': task_success_alert,
    'on_failure_callback': task_failure_alert
}

This will allow you to customize before you call the function in your package.

On a side note, airflow has smtp email functionality. Instead of writing your own solution, you can utilize those.

Oscar Foley
  • 6,817
  • 8
  • 57
  • 90
nightgaunt
  • 890
  • 12
  • 30
  • When I use CONTEXT it my tasks finishes but dag keeps running.It never ends. any suggestions. – Akshay Lande Jun 21 '20 at 06:35
  • Can't understand or debug without a code example. Add a question and tag it airflow. I will check and answer if I am able to – nightgaunt Jun 22 '20 at 04:10
  • def notify_email(context): import inspect """Send custom email alerts.""" import smtplib, ssl from email.mime.text import MIMEText from email.mime.multipart import MIMEMultipart sender_email = 'abc@gmail.com' receiver_email = 'xyz@gmail.com' password = "abc" message = MIMEMultipart("alternative") #task_instance = context['task'].task_id dag_instance=context['dag_id'].dag_id When using context['dag_id'].dag_id my dag keeps on running task finishes and mail is not sent – Akshay Lande Jun 22 '20 at 09:19
  • @Gabriel Eckers The SMTP util in airflow still provides an option to send emails as a Task. While on_success option is not available, email functionality can be kept within airflow DAG. – nightgaunt Jun 23 '20 at 04:59
5

You can use partial to create a function with a predefined argument like:

from functools import partial
new_task_success_alert = partial(task_success_alert, email='your_email')

And then add the new function as a callback:

on_success_callback=new_task_success_alert
Franco Piccolo
  • 6,845
  • 8
  • 34
  • 52
  • This answer is much better than the highlighted response because the former relies on knowing the file name at the time of import. – Javier Muñoz Aug 29 '23 at 15:53
3

You can create a task that its only purpose is to push configuration setting through xcoms. You can pull the configuration via context as the task_instance object is included in context.

def push_configuration(ti, params):
    ti.xcom_push(key='conn_id', value=params)

def task_success_alert(context):
    ti = context.get('ti') 
    params = ti.xcom_pull(key='params', task_ids='Settings')
    ...


step0 = PythonOperator(
        task_id='Settings',
        python_callable=push_configuration,
        op_kwargs={'params': params})

step1 = BashOperator(
        task_id='step1',
        bash_command='pwd',
        on_success_callback=task_success_alert)
kemazariegos
  • 41
  • 2