读取CSV并加载到PostgreSQL的气流管道

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime
import pandas as pd
import psycopg2

# 定义默认参数
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2023, 1, 1),
}

# 定义DAG
dag = DAG('csv_to_postgres', default_args=default_args, schedule_interval='@daily')

def load_csv_to_postgres():
    # 读取CSV文件
    df = pd.read_csv('/path/to/your/file.csv')
    
    # 连接到PostgreSQL数据库
    conn = psycopg2.connect(
        host="localhost",
        database="your_database",
        user="your_user",
        password="your_password"
    )
    
    # 创建游标
    cur = conn.cursor()
    
    # 插入数据
    for index, row in df.iterrows():
        cur.execute(
            "INSERT INTO your_table (column1, column2, column3) VALUES (%s, %s, %s);",
            (row['column1'], row['column2'], row['column3'])
        )
    
    # 提交事务
    conn.commit()
    
    # 关闭连接
    cur.close()
    conn.close()

# 定义任务
load_task = PythonOperator(
    task_id='load_csv_to_postgres',
    python_callable=load_csv_to_postgres,
    dag=dag,
)

# 设置任务依赖
load_task

if __name__ == "__main__":
    dag.cli()