如何从Python函数向Airflow Athena Operator传递参数？

方法1：直接用Airflow模板变量引用dag_run配置 #

AWSAthenaOperator原生支持Jinja2模板语法，可直接从dag_run.conf里提取运行时参数，无需额外函数处理：

from airflow.providers.amazon.aws.operators.athena import AWSAthenaOperator

athena_query_task = AWSAthenaOperator(
    task_id='run_athena_query',
    query="SELECT * FROM {{ dag_run.conf['base'] }}.{{ dag_run.conf['tabela'] }}",
    database="{{ dag_run.conf['base'] }}",
    output_location="s3://your-bucket/path/{{ dag_run.conf['base'] }}/{{ dag_run.conf['tabela'] }}/",
    aws_conn_id='aws_default',
    dag=dag
)

这种方式适合参数逻辑简单的场景，直接通过模板语法完成参数注入。

方法2：Python函数生成参数+XCom传递 #

如果需要调用自定义ETL类处理复杂逻辑生成参数，可以用PythonOperator把参数存入XCom，再让AWSAthenaOperator引用：

步骤1：编写参数生成函数 #

def generate_athena_params(**context):
    # 从dag_run获取配置
    base = context['dag_run'].conf['base']
    tabela = context['dag_run'].conf['tabela']
    
    # 调用自定义ETL类生成参数
    etl = YourCustomETLClass()
    query = etl.get_query(base, tabela)
    output_location = etl.get_output_location(base, tabela)
    database = base
    
    # 存入XCom供后续任务使用
    ti = context['ti']
    ti.xcom_push(key='athena_query', value=query)
    ti.xcom_push(key='athena_output', value=output_location)
    ti.xcom_push(key='athena_db', value=database)

步骤2：串联任务 #

from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.operators.athena import AWSAthenaOperator

generate_params_task = PythonOperator(
    task_id='generate_athena_params',
    python_callable=generate_athena_params,
    provide_context=True,
    dag=dag
)

athena_query_task = AWSAthenaOperator(
    task_id='run_athena_query',
    query="{{ ti.xcom_pull(task_ids='generate_athena_params', key='athena_query') }}",
    database="{{ ti.xcom_pull(task_ids='generate_athena_params', key='athena_db') }}",
    output_location="{{ ti.xcom_pull(task_ids='generate_athena_params', key='athena_output') }}",
    aws_conn_id='aws_default',
    dag=dag
)

generate_params_task >> athena_query_task

方法3：TaskFlow API简化参数传递（Airflow 2.0+） #

如果用Airflow 2.x，TaskFlow API会自动处理XCom传递，代码更简洁：

from airflow.decorators import dag, task
from airflow.providers.amazon.aws.operators.athena import AWSAthenaOperator
from datetime import datetime

@dag(start_date=datetime(2023, 1, 1), schedule_interval=None)
def athena_etl_dag():
    @task
    def generate_athena_params(dag_run=None):
        base = dag_run.conf['base']
        tabela = dag_run.conf['tabela']
        
        etl = YourCustomETLClass()
        return {
            'query': etl.get_query(base, tabela),
            'database': base,
            'output_location': etl.get_output_location(base, tabela)
        }
    
    params = generate_athena_params()
    
    athena_query_task = AWSAthenaOperator(
        task_id='run_athena_query',
        query=params['query'],
        database=params['database'],
        output_location=params['output_location'],
        aws_conn_id='aws_default'
    )
    
    params >> athena_query_task

athena_etl_dag()

注意事项 #

• 若传递大SQL查询，注意Airflow默认XCom大小限制（48KB），超出可将查询存S3再引用路径，或调整XCom配置；
• 测试时通过Airflow UI触发DAG，传入{"base": "你的数据库名", "tabela": "你的表名"}验证参数是否生效。

内容的提问来源于stack exchange，提问作者Fábio Elias Reis Ritter