【Airflow】基于数据的调度 -条件触发

条件触发更新

任务解析

工具特性

Airflow的数据调度机制是以上游的producer DAG的结果为准，触发下游的Consumer DAG的。也就是说，producer中的任务执行的结果可以看作consumer的触发条件。producer执行成功触发consumer，producer失败则不执行consumer。所以consumer中不能自身控制是否进行触发，触发条件的逻辑实现在于producer的执行。

所以数据调度的控制条件可以放在producer的最后一个任务中，或者是consumer的第一个任务中。但是放在producer任务中能够相对更加节约资源。

代码

import os
import time
import gzip
import random
import pandas as pd
import numpy as np
from shutil import copy
import sqlite3
import json
import xml.etree.ElementTree as ET
from datetime import datetime,timedelta

from airflow import DAG
from airflow.models import Variable
from airflow.datasets import Dataset
from airflow.decorators import dag,task
from datetime import datetime,timedelta
from airflow.models.baseoperator import chain
from airflow.operators.python import PythonOperator
from airflow.models import Variable

TEST_PATH = '/home/xxxxx/airflow/test_demo'
BUFFER_PATH = '/home/xxxxx/airflow/test_demo/buffer'
FILE_PATH = os.path.join(TEST_PATH,'files')
RESULT_PATH = os.path.join(TEST_PATH,'result')
DATA_PATH = '/home/xxxxx/airflow/dataset/xxxxx'
buffer_dataset= Dataset(BUFFER_PATH)
test_dataset= Dataset(TEST_PATH)

def xml_to_table(input_data_path):
    '''
    transfor xml data into a dataframe
    '''
    with gzip.open(input_data_path,'rb') as f:
        content = f.read()
        df = pd.read_xml(content)
        df = df[3::]
        df.to_csv(os.path.join(RESULT_PATH,input_data_path.split('/')[-1]+'.csv'),index=False)
    return df,os.path.join(RESULT_PATH,input_data_path.split('/')[-1]+'.csv')

#! data source
with DAG(
    dag_id = 'fake_move',
    tags = ['test','demo'],
    start_date=datetime(2024,7,17),
    schedule = timedelta(minutes=5)
)as fake_move:
    def move():
        if not os.path.exists(TEST_PATH):
            os.makedirs(TEST_PATH,exist_ok=True)
        if not os.path.exists(RESULT_PATH):
            os.makedirs(RESULT_PATH,exist_ok=True)
        if not os.path.exists(BUFFER_PATH):
            os.makedirs(BUFFER_PATH,exist_ok=True)
        if not os.path.exists(FILE_PATH):
            os.makedirs(FILE_PATH,exist_ok=True)

        upstream_files = os.listdir(DATA_PATH)
        with open('/home/xxxxx/airflow/test_index.txt','r') as index:
            INDEX = int(index.read())
        src = os.path.join(DATA_PATH,upstream_files[INDEX])
        print(f'move NO.{INDEX} file which name is :',upstream_files[INDEX])
        dst = os.path.join(BUFFER_PATH,upstream_files[INDEX])
        copy(src,dst)
        with open('/home/xxxxx/airflow/test_index.txt','w') as index:
            INDEX-=1
            print(INDEX,file=index)

    move_task = PythonOperator(
        task_id='move',
        python_callable=move
    )

#! get data from the buffer(which is the folder that publish the data files)
with DAG(
    dag_id='conditional_get_xml',
    tags=['data','test','demo'],
    start_date=datetime(2024,7,17),
    schedule=timedelta(minutes=2),
    max_active_runs=1,
    
)as dag_update:
    
    def get_target():
        if not os.path.exists(TEST_PATH):
            os.makedirs(TEST_PATH,exist_ok=True)
        if not os.path.exists(RESULT_PATH):
            os.makedirs(RESULT_PATH,exist_ok=True)
        if not os.path.exists(BUFFER_PATH):
            os.makedirs(BUFFER_PATH,exist_ok=True)
        if not os.path.exists(FILE_PATH):
            os.makedirs(FILE_PATH,exist_ok=True)

        xml_files = os.listdir(BUFFER_PATH)
        for xml in xml_files:
            if xml.startswith('._'):
                os.remove(os.path.join(BUFFER_PATH,xml))
        upstream_files = os.listdir(BUFFER_PATH)
        downstream_files = os.listdir(FILE_PATH)
        target_files = list(set(upstream_files)-set(downstream_files))

        #move files
        for target in target_files:
            src = os.path.join(BUFFER_PATH,target)
            dst = os.path.join(FILE_PATH,target)
            copy(src,dst)
        return target_files

    def update_condition(condition_amount=5000,**kwargs):
        ''' 
        definte the update rule
        '''
        ti = kwargs['ti']
        file_list = ti.xcom_pull(task_ids="get_target")
        amount = 0
        reasult_df = []
        file_names = []
        for file in file_list:
            file,file_name = xml_to_table(os.path.join(FILE_PATH,file))
            reasult_df.append(file)
            file_names.append(file_name)
            amount += len(file)
        with open(os.path.join(TEST_PATH,'record.json'),'w') as recorder:
            json.dump(file_names,recorder,ensure_ascii=False)

        print('amount:',amount)
        flag = True if amount>condition_amount else False
        if flag:
            return 
        else:
            flag = 243 / 0
            
    get_target_task = PythonOperator(
        task_id = 'get_target',
        outlets = test_dataset,
        python_callable=get_target
    )

    condition_task = PythonOperator(
        task_id = 'condition_task',
        python_callable=update_condition
    )
    chain(get_target_task,condition_task)

#! deal with the data
with DAG(
    dag_id='conditional_data_deal',
    tags=['data','test','demo'],
    start_date=datetime(2024,7,17),
    schedule=[test_dataset]
)as dag_data_aware:
    
    def output():
        con = sqlite3.connect(os.path.join(TEST_PATH,'xxxxx.db'))
        file_list = os.listdir(FILE_PATH)
        file_list.sort(key=lambda fn:os.path.getmtime(FILE_PATH+'/'+fn))
        with open(os.path.join(TEST_PATH,'record.json'),'r') as recorder:
            result_files = json.load(recorder)
        for file in result_files:
            df = pd.read_csv(os.path.join(RESULT_PATH,file))
            df.to_sql(name='xxxxx',con = con,if_exists='append',index=False)
        con.close()

    output_task = PythonOperator(
        task_id = 'output_task',
        python_callable = output
    )

    chain(output_task)

if __name__=='__main__':
    fake_move.test()
    dag_update.test()
    # dag_data_aware.test()

任务划分

• fake_move：模拟数据的产生
  • move：将数据从datasets搬运至buffer文件夹（buffer即一个约定的数据获取文件夹）
• producer
  • get_target_task：得出与上游相比，下游缺少的文件名
  • update_condtion：控制下游的任务是否被触发的任务，在此加入想要加入的触发条件即可
• consumer：
  • output_task : 对于处理后的数据进行入库或者更新操作

文件夹划分

DATA_PATH：存储原始数据的文件夹

BUFFER_PATH：模拟上游被同步的文件夹

FILE_PATH：consumer监控的文件夹，也就是目标文件夹

方案思路

方案一

参照实现思路中的描述，故意在producer的最后一个任务中不满足条件时则将任务设为失败

结果：上游的任务失败后可以实现不触发新的任务

方案二

producer中最后一个condition task中的逻辑改为：如果不满足更新条件，则不会执行将文件同步至FLIE_PATH的操作（即不会触发FLIE_PATH的变化）。

结果：保证目标文件夹不被修改，就可以不触发consumer任务

结论

实现条件触发的方式有三种

1. 一种是保证producer任务中与下游对接的任务存在控制条件（和官方文档的描述最为符合，并且有助于运维）
2. 保证目标文件夹没有修改（写起来对于其他部分的逻辑可能影响比较多）
3. 在consumer任务中设置控制条件（比较别扭的逻辑，可能逼不得已了可以使用）