第12章 数据管道¶

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

# 定义DAG
default_args = {
    'owner': 'data_team',
    'start_date': datetime(2024, 1, 1),
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'depends_on_past': False,
}

dag = DAG(
    'sales_pipeline',
    default_args=default_args,
    description='Daily sales ETL pipeline',
    schedule='@daily',
    catchup=False,
    tags=['sales', 'etl']
)

# 定义任务
def extract_sales():
    """抽取销售数据"""
    print("Extracting sales data...")
    # 实际的数据抽取逻辑
    return 'sales_data.csv'

def transform_sales(ti):
    """转换销售数据"""
    print("Transforming sales data...")
    # 实际的数据转换逻辑
    return 'transformed_sales.csv'

def load_sales(ti):
    """加载销售数据"""
    print("Loading sales data...")
    # 实际的数据加载逻辑
    return 'load_complete'

# 创建任务
extract_task = PythonOperator(
    task_id='extract',
    python_callable=extract_sales,
    dag=dag
)

transform_task = PythonOperator(
    task_id='transform',
    python_callable=transform_sales,
    dag=dag
)

load_task = PythonOperator(
    task_id='load',
    python_callable=load_sales,
    dag=dag
)

# 设置任务依赖
extract_task >> transform_task >> load_task

from apscheduler.schedulers.background import BackgroundScheduler
from datetime import datetime

class TaskScheduler:
    """任务调度器"""

    def __init__(self):
        self.scheduler = BackgroundScheduler()
        self.jobs = {}

    def add_job(self, job_id, func, trigger, **kwargs):  # *args接收任意位置参数；**kwargs接收任意关键字参数
        """添加任务"""
        self.scheduler.add_job(func, trigger, id=job_id, **kwargs)
        self.jobs[job_id] = func
        print(f"Job added: {job_id}")

    def remove_job(self, job_id):
        """移除任务"""
        self.scheduler.remove_job(job_id)
        if job_id in self.jobs:
            del self.jobs[job_id]
        print(f"Job removed: {job_id}")

    def list_jobs(self):
        """列出任务"""
        jobs = self.scheduler.get_jobs()
        for job in jobs:
            print(f"Job: {job.id}, Next run: {job.next_run_time}")

    def start(self):
        """启动调度器"""
        self.scheduler.start()
        print("Scheduler started")

    def shutdown(self):
        """关闭调度器"""
        self.scheduler.shutdown()
        print("Scheduler shutdown")

# 使用示例
scheduler = TaskScheduler()

def extract_data():
    """抽取数据"""
    print("Extracting data...")

def transform_data():
    """转换数据"""
    print("Transforming data...")

def load_data():
    """加载数据"""
    print("Loading data...")

# 添加任务
scheduler.add_job('extract', extract_data, 'interval', seconds=3600)
scheduler.add_job('transform', transform_data, 'interval', seconds=3600)
scheduler.add_job('load', load_data, 'interval', seconds=3600)

# 启动调度器
scheduler.start()

# 列出任务
scheduler.list_jobs()