Airflow 创建动态 FileSensorAsync 任务

Question

我正在将 Airflow 2.3.3 与 Astronomer 自定义模块一起使用。 我的目标是建立一个 dag

1. 异步轮询状态 json
2. 加载 json 并继续基于 json 内容的特定分支
3. 根据 json 变量值创建 X 个异步轮询任务

我尝试了几种方法，但不知何故，我无法让动态任务创建工作如我所愿。

如果有人可以查看我的 dag 文件并建议我如何正确实施或指出我的方法中的缺陷，我将不胜感激。

                                                                                                                                                                                                                      
    with DAG(                                                                                                                                                                                                                                                                                                 
        dag_id="dynamic-dag-example",                                                                                                                                                                                                                                                                   
        start_date=datetime(2022, 1, 1),                                                                                                                                                                                                                                                                      
        schedule_interval=None,                                                                                                                                                                                                                                                                               
        #schedule_interval='*/5 * * * *',                                                                                                                                                                                                                                                                     
        dagrun_timeout=timedelta(minutes=6),                                                                                                                                                                                                                                                                  
        catchup=False,                                                                                                                                                                                                                                                                                        
        default_args=default_args,                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                
    ) as dag:                                                                                                                                                                                                                                                                                                 
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
        def choose_branch_overall_status(**context):                                                                                                                                                                                                                                                          
            #overall_status = load_site_status_json(kwargs['site'], kwargs['timestamp'])['overall_status']                                                                                                                                                                                                    
            overall_status = context['templates_dict']['overall_status']                                                                                                                                                                                                                                      
            if overall_status == 'ok':                                                                                                                                                                                                                                                                        
                return ['overall_status_ok']                                                                                                                                                                                                                                                                  
            if overall_status == 'critical':                                                                                                                                                                                                                                                                  
                return ['overall_status_critical']                                                                                                                                                                                                                                                            
            return ['overall_status_undefined']                                                                                                                                                                                                                                                               
                                                                                                                                                                                                                                                                                                              
        for site in ['site1']:                                                                                                                                                                                                                                                                                
            file_sensor_task = FileSensorAsync(                                                                                                                                                                                                                                                               
                task_id=f"poll_status_async_{site}",                                                                                                                                                                                                                                                          
                filepath=f"{{{{ data_interval_start.strftime('%Y%m%d%H%M') }}}}_{ site }.json",                                                                                                                                                                                                               
                fs_conn_id=FS_CONN_ID                                                                                                                                                                                                                                                                         
            )                                                                                                                                                                                                                                                                                                 
                                                                                                                                                                                                                                                                                                              
            @task                                                                                                                                                                                                                                                                                             
            def load_site_status_json(**kwargs) -> Dict:                                                                                                                                                                                                                                                      
                fs_hook = FSHook(conn_id=FS_CONN_ID)                                                                                                                                                                                                                                                          
                file_path = os.path.join(fs_hook.get_path(), kwargs['data_interval_start'].strftime('%Y%m%d%H%M') + f'_{ site }.json')                                                                                                                                                                        
                #file_path = os.path.join(fs_hook.get_path(), timestamp + f'_{ site }.json')                                                                                                                                                                                                                  
                with open(file_path, 'r') as json_file:                                                                                                                                                                                                                                                       
                    status_json = json.load(json_file)                                                                                                                                                                                                                                                        
                    return status_json                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               
                                                                                                                     
            check_overall_status = BranchPythonOperator(                                                                                                                                                                                                                                                      
                task_id = 'check_overall_status',                                                                                                                                                                                                                                                             
                python_callable = choose_branch_overall_status,                                                                                                                                                                                                                                               
                provide_context = True,                                                                                                                                                                                                                                                                                                                                                                                                                                                                          
                templates_dict={'overall_status': "{{ ti.xcom_pull(task_ids='extract')['overall_status'] }}",                                                                                                                                                                                                 
                                'angles': "{{ ti.xcom_pull(task_ids='extract')['angles'] }}"})                                                                                                                                                                                                                
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              
            overall_status_ok = DummyOperator(task_id='overall_status_ok')                                                                                                                                                                                                                                    
            overall_status_critical = DummyOperator(task_id='overall_status_critical')                                                                                                                                                                                                                        
            overall_status_undefined = DummyOperator(task_id='overall_status_undefined')                                                                                                                                                                                                                      
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                   
            @task                                                                                                                                                                                                                                                                                             
            def load_angles(**kwargs) -> List:                                                                                                                                                                                                                                                                
                """Instead of this task, read angles from the xcom of load_site_status_json                                                                                                                                                                                                                   
                """                                                                                                                                                                                                                                                                                           
                return ['A','B','C']                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                             
                                                                                                                                                                                                                                                                                                              
            @task(trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)                                                                                                                                                                                                                                       
            def print_angle(angle):                                                                                                                                                                                                                                                                           
                print(angle)                                                                                                                                                                                                                                                                                  
                                                                                                                                                                                                                                                                                                              
            file_sensor_task >> load_site_status_json() >> check_overall_status >> [overall_status_ok, overall_status_critical, overall_status_undefined]                                                                                                                                                                                                                                                                                                                                                                                                            
            overall_status_ok >> print_angle.expand(angle=load_angles())                                            

            #for angle in load_site_status_json['angles']:                                                                                                                                                                                                                                                    
            #    - Create AsyncFileSensor that polls for file timestamp_site_angle                                                                                                                                                                                                                

编辑：我找到了一个使用部分和扩展的解决方案，但如果有“for loop”解决方案也很好，这样我就可以继续分别处理每个文件。

Answer 1

花了几天时间才把它弄好，但显然神奇的是在这里问，这样我的大脑就能弄清楚。 在这里，我希望它在未来对其他人有所帮助。

        @task(trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)                                                                                                                                                                                                                                       
        def render_filenames():                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
            # TODO: add error handling when theres no angles etc.                                                                                                                                                                                                                                         
            context = get_current_context()                                                                                                                                                                                                                                                               
            timestamp = context['data_interval_start'].strftime('%Y%m%d%H%M')                                                                                                                                                                                                                             
            angles = context['ti'].xcom_pull(task_ids='load_site_status_json')['angles']                                                                                                                                                                                                                  
            return [f"{ timestamp }_{ site }_{ angle }.txt" for angle in angles]                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                
        angle_sensor_tasks = FileSensorAsync.partial(task_id="fsa", fs_conn_id=FS_CONN_ID, trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS).expand(                                                                                                                                                  
            filepath=render_filenames()                                                                                                                                                                                                                                                                   
        )                                                                                                                                                                                                                                                                                                 
                                                                                                                                                                                                                                                                                                          
        overall_status_ok >> angle_sensor_tasks >> DummyOperator(task_id="END")