Airflow: Not able to transfer data from myql database to csv file

Viewed 1121

I'm working on Airflow where I'm trying to transfer data form mysql database to csv file. Below are the codes and functions

from airflow import DAG
from datetime import datetime,timedelta
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow.operators.mysql_operator import MySqlOperator

from dbextract import extract_data
from dbpost_processing import dbpost_process

default_args = {"owner":"airflow","start_date":datetime(2021,7,10)}
with DAG(dag_id="dbworkflow2",default_args=default_args,schedule_interval=None) as dag:
   
    extract = MySqlOperator(
        task_id='extract',
        mysql_conn_id="mysql_db1", 
        sql = extract_data
        )
        
    dbpost_process = PythonOperator(
        task_id = "dbpost_process",
        python_callable = dbpost_process
        )   
        
    extract >> dbpost_process
import pandas as pd
def extract_data():
    df=pd.read_sql('SELECT * FROM new_table', mysql_conn_id)
import pandas as pd
def dbpost_process():
    df.to_csv('~/op_files/sample3.csv', index=False)

Getting error as function object is not iterable in extract_data step

1 Answers

The sql in MySQL file is not a callable, it should be either string or list of strings. You are also trying to pass panda dataframe between tasks and it will not work because the tasks might (and likely will) run on a different machine within different processes.

The way how the data is exchanged between tasks usuallu is via XComs (for small amount of data it should be via DB, for bigger amounts of data you can add custom XCom backend and pass the data for example via S3 or GCS).

However in your case you do not need to have two separate tasks/operators. instead you should use MySQL Hook inside the Python operator to read the data and process it in the same tasks. Splitting that job into two separate tasks makes no sense - MySQL Operator is really to execute DDL or DML operations, not to extract the data (precisely because Airflow Operators work in isolation)

Airflow Has instead the concept of Hooks that provide the API you can use to run your query and process data within the same Python Operator callable. And even recently it can all be done using @task decorator so it is super simple and much easier to write - especially if you are used to writing functional Python.

See https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html

But in your case you don't need to do even that because you want to use pandas to communicate with MySQL db, so you won't even need to use the Hook. Something like that is enough. You shoudl not use conn_id in this case, but you need to pass an SQL-Alchemy compatible connection string there. Not sure if the Airflow connection url will work but if it does, then you can use Connection.get_uri() (maybe you will have to adapt the URI a bit).

Something like that (this is an inspiration for you, not compilable code, so you need to work out the details) should work:

@dag(default_args=default_args, schedule_interval=None, start_date=days_ago(2))
def my_dag():
    @task()
    def my_extraction():
         df =  pd.read_sql('SELECT *', Connnection.get("my_connection_id").get_uri()) 
         file = post_process(df)

Now - you will have to do something with that file as well, because once the task is complete the local data will not be available (unless you use LocalExecutor) - you can send the CSV somewhere (and for that you can use any Hook - for example S3Hook).

Related