Python脚本数据管道如何用Airflow编排

本文目录导读：

1. Airflow基础架构
2. Python脚本数据管道示例
3. Airflow DAG定义
4. 高级编排模式
5. 监控和告警
6. 最佳实践
7. 部署和运维

我将详细介绍如何使用Airflow编排Python脚本数据管道。

Airflow基础架构

安装Airflow

# 安装Airflow
pip install apache-airflow
# 初始化数据库
airflow db init
# 创建管理员用户
airflow users create \
    --username admin \
    --firstname Admin \
    --lastname User \
    --role Admin \
    --email admin@example.com \
    --password admin

Python脚本数据管道示例

创建Python脚本 (scripts/)

# scripts/data_extract.py
import pandas as pd
from sqlalchemy import create_engine
def extract_data(source_config):
    """数据抽取函数"""
    engine = create_engine(source_config['connection_string'])
    query = """
    SELECT * FROM source_table 
    WHERE date >= CURRENT_DATE - INTERVAL '1 day'
    """
    df = pd.read_sql(query, engine)
    df.to_parquet('/tmp/raw_data.parquet')
    return '/tmp/raw_data.parquet'
# scripts/data_transform.py
def transform_data(input_path, output_path):
    """数据转换函数"""
    df = pd.read_parquet(input_path)
    # 数据清洗和转换
    df['clean_column'] = df['raw_column'].str.strip().lower()
    df = df.dropna(subset=['important_field'])
    df['processed_date'] = pd.Timestamp.now()
    # 聚合操作
    aggregated = df.groupby('category').agg({
        'value': ['sum', 'mean', 'count']
    })
    aggregated.to_parquet(output_path)
    return output_path
# scripts/data_load.py
def load_data(input_path, target_config):
    """数据加载函数"""
    engine = create_engine(target_config['connection_string'])
    df = pd.read_parquet(input_path)
    df.to_sql(
        'target_table',
        engine,
        if_exists='replace',
        index=False,
        method='multi',
        chunksize=1000
    )
    return f"Loaded {len(df)} rows to target table"

Airflow DAG定义

创建DAG文件 (dags/data_pipeline_dag.py)

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.dummy import DummyOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
import sys
sys.path.append('/path/to/scripts')
from data_extract import extract_data
from data_transform import transform_data
from data_load import load_data
# 默认参数
default_args = {
    'owner': 'data_team',
    'depends_on_past': False,
    'email': ['alerts@example.com'],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'start_date': datetime(2024, 1, 1),
}
# 定义DAG
dag = DAG(
    'python_data_pipeline',
    default_args=default_args,
    description='Python脚本数据管道编排',
    schedule_interval='0 2 * * *',  # 每天凌晨2点执行
    catchup=False,
    tags=['data_pipeline'],
)
# 定义任务函数
def extract_task(**context):
    """抽取任务"""
    source_config = {
        'connection_string': 'postgresql://user:pass@host:5432/source_db'
    }
    try:
        file_path = extract_data(source_config)
        context['task_instance'].xcom_push(key='raw_file_path', value=file_path)
        return file_path
    except Exception as e:
        raise AirflowException(f"抽取失败: {str(e)}")
def transform_task(**context):
    """转换任务"""
    ti = context['task_instance']
    input_path = ti.xcom_pull(task_ids='extract_data', key='raw_file_path')
    output_path = '/tmp/transformed_data.parquet'
    try:
        result_path = transform_data(input_path, output_path)
        ti.xcom_push(key='transformed_file_path', value=result_path)
        return result_path
    except Exception as e:
        raise AirflowException(f"转换失败: {str(e)}")
def load_task(**context):
    """加载任务"""
    ti = context['task_instance']
    input_path = ti.xcom_pull(task_ids='transform_data', key='transformed_file_path')
    target_config = {
        'connection_string': 'postgresql://user:pass@host:5432/target_db'
    }
    try:
        result = load_data(input_path, target_config)
        print(result)
        return result
    except Exception as e:
        raise AirflowException(f"加载失败: {str(e)}")
# 创建任务
start_pipeline = DummyOperator(
    task_id='start_pipeline',
    dag=dag,
)
extract = PythonOperator(
    task_id='extract_data',
    python_callable=extract_task,
    provide_context=True,
    dag=dag,
)
transform = PythonOperator(
    task_id='transform_data',
    python_callable=transform_task,
    provide_context=True,
    dag=dag,
)
load = PythonOperator(
    task_id='load_data',
    python_callable=load_task,
    provide_context=True,
    dag=dag,
)
# 数据质量检查
quality_check = PostgresOperator(
    task_id='quality_check',
    postgres_conn_id='target_db',
    sql="""
        SELECT COUNT(*) as row_count 
        FROM target_table 
        WHERE processed_date = CURRENT_DATE
    """,
    dag=dag,
)
# 发送通知
def send_notification(**context):
    """发送任务通知"""
    import smtplib
    from email.message import EmailMessage
    ti = context['task_instance']
    status = context['task_instance'].state
    msg = EmailMessage()
    msg.set_content(f"数据管道状态: {status}\n执行时间: {datetime.now()}")
    msg['Subject'] = f"数据管道执行通知 - {status}"
    msg['To'] = 'team@example.com'
    # 发送邮件
notification = PythonOperator(
    task_id='send_notification',
    python_callable=send_notification,
    provide_context=True,
    dag=dag,
)
end_pipeline = DummyOperator(
    task_id='end_pipeline',
    dag=dag,
)
# 设置任务依赖
start_pipeline >> extract >> transform >> load >> quality_check >> notification >> end_pipeline

高级编排模式

分支和条件执行

from airflow.operators.python import BranchPythonOperator
def check_data_quality(**context):
    """数据质量检查分支"""
    # 检查数据质量
    if quality_score > 0.95:
        return 'continue_pipeline'
    else:
        return 'send_alert'
branch = BranchPythonOperator(
    task_id='quality_branch',
    python_callable=check_data_quality,
    provide_context=True,
    dag=dag,
)
# 并行处理
def process_chunk(chunk_id, **context):
    """并行处理数据块"""
    # 处理逻辑
parallel_tasks = []
for i in range(5):
    task = PythonOperator(
        task_id=f'process_chunk_{i}',
        python_callable=process_chunk,
        op_kwargs={'chunk_id': i},
        provide_context=True,
        dag=dag,
    )
    parallel_tasks.append(task)
# 并行执行
start_task >> parallel_tasks >> merge_task

监控和告警

配置监控

# alert_config.py
from airflow.models import Variable
from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator
SLACK_CONN_ID = 'slack_default'
def task_failure_alert(context):
    """任务失败告警"""
    return SlackWebhookOperator(
        task_id='slack_alert',
        slack_webhook_conn_id=SLACK_CONN_ID,
        message=f"""
        :red_circle: **数据管道失败**
        DAG: {context['dag'].dag_id}
        Task: {context['task'].task_id}
        时间: {context['execution_date']}
        """,
        dag=dag,
    ).execute(context=context)

最佳实践

配置管理

# airflow.cfg
[core]
executor = LocalExecutor
parallelism = 32
dag_concurrency = 16
[scheduler]
dag_dir_list_interval = 30
min_file_process_interval = 30
[smtp]
smtp_host = smtp.gmail.com
smtp_starttls = True
smtp_ssl = False
smtp_user = your_email@gmail.com
smtp_password = your_password
smtp_port = 587
smtp_mail_from = your_email@gmail.com

资源管理

# 使用资源池
from airflow.operators.python import PythonOperator
task = PythonOperator(
    task_id='heavy_computation',
    python_callable=heavy_function,
    pool='data_processing_pool',
    priority_weight=10,
    pool_slots=2,
    dag=dag,
)

部署和运维

Docker部署

# Dockerfile
FROM apache/airflow:2.7.0
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY dags/ /opt/airflow/dags/
COPY scripts/ /opt/airflow/scripts/
COPY config/ /opt/airflow/config/

docker-compose.yml

version: '3.8'
services:
  postgres:
    image: postgres:13
    environment:
      POSTGRES_USER: airflow
      POSTGRES_PASSWORD: airflow
      POSTGRES_DB: airflow
  webserver:
    build: .
    depends_on:
      - postgres
    ports:
      - "8080:8080"
    volumes:
      - ./dags:/opt/airflow/dags
      - ./scripts:/opt/airflow/scripts

这个架构提供了完整的Python数据管道编排方案,包括数据抽取、转换、加载的完整流程，以及监控、告警、错误处理等企业级功能。