airflow_presto_operator_dag_with_pandas_handler.py

1import pandas as pd
2from datetime import datetime
3
4from airflow import DAG
5from airflow.providers.presto.operators.presto import PrestoOperator
6
7with DAG(
8    dag_id="example_presto",
9    schedule=None,
10    start_date=datetime(2022, 1, 1),
11    catchup=False,
12    tags=["example"],
13) as dag:
14    # [START howto_operator_presto]
15    handler = PrestoOperator(
16        task_id="select_presto_query",
17        sql="SELECT * FROM table",
18        handler=list,
19    )
20    # [END howto_operator_presto]
21
22    # [START howto_operator_presto_multiple_queries]
23    multiple_queries = PrestoOperator(
24        task_id="multiple_queries",
25        sql=["SELECT * FROM table1", "SELECT * FROM table2"],
26    )
27    # [END howto_operator_presto_multiple_queries]
28
29    # [START howto_operator_presto_pandas_handler]
30    def pandas_handler(cursor, **kwargs):
31        return pd.DataFrame(cursor.fetchall())
32
33    pandas_query = PrestoOperator(
34        task_id="pandas_query",
35        sql="SELECT * FROM table",
36        handler=pandas_handler,
37    )
38    # [END howto_operator_presto_pandas_handler]