Airflow - Attach file to EmailOperator from pandas

user_18213714 0
Viewed 317

I got pandas dataframe from sql server and put it in variable. How can i attach this dataframe as excel-file to EmailOperator in airflow (w/o downloading it to the local system)? My code without putting into file)):

sql_select = 'SELECT TOP 100 * FROM test_table'
df = pd.read_sql(sql=sql_select, con=sql_engine)
????

email = EmailOperator(task_id="email_task",
                          to="sample@sample.com",
                          subject="Sample_subject",
                          html_content="Sample_message",
                          files=???,
                          dag=dag)
airflow python
1 Answers

Operators can not take dataframe as input. Operators do not share in-memory objects between them.

I would recommend you to create a custom operator: SqlToEmailOperator The operator will run the query, store the output as csv in local disk and then send the file to email. You didn't mention which DB you are using so for the example I would assume it's MySQL.

A skeleton solution would be something like:

from tempfile import NamedTemporaryFile

from airflow.models import BaseOperator
from airflow.providers.mysql.hooks.mysql import MySqlHook
from airflow.utils.email import send_email


class MySQLToEmailOperator(BaseOperator):

    def __init__(
        self,
        *,
        mysql_conn_id,
        sql,
        to,
        subject,
        html_content,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.mysql_conn_id = mysql_conn_id,
        self.sql = sql,
        self.to = to,
        self.subject = subject,
        self.html_content = html_content

    def execute(self):
        hook = MySqlHook(mysql_conn_id=self.mysql_conn_id)
        df = hook.get_pandas_df(self.sql)
        with NamedTemporaryFile() as tmp_file:
            df.to_csv(tmp_file.name)
            send_email(
                self.to,
                self.subject,
                self.html_content,
                files=tmp_file.name,
            )

I didn't test it but that is the general idea. then you can use this operator as:

MySQLToEmailOperator(task_id="email_task",
                     mysql_conn_id="sql_engine",
                     sql="SELECT TOP 100 * FROM test_table",
                     to="sample@sample.com",
                     subject="Sample_subject",
                     html_content="Sample_message",
                     dag=dag
                     )
user_146244090
Related