参数中的 Airflow Jinja 模板

Question

我有一个 Airflow 运算符，它允许我查询 Athena，它接受 Jinja 模板文件作为查询输入。通常，我将表/数据库名称等变量传递给模板以创建表并添加分区语句。这对于定义的字符串效果很好。

我的任务定义如下：

        db = 'sample_db'
        table = 'sample_table'
        out = 's3://sample'
        p1='2020'
        p2='1'

        add_partition_task= AWSAthenaOperator(
            task_id='add_partition',
            query='add_partition.sql',
            params={'database': db,
                    'table_name': table,
                    'p1': p1
                    'p2': p2},
            database=db,
            output_location=out
        )

模板化的 SQL 文件如下所示：

ALTER TABLE {{ params.database }}.{{ params.table_name }} ADD IF NOT EXISTS
PARTITION (partition1= '{{ params.p1 }}', partition2 = '{{ params.p2 }}')

这个模板效果很好。

对此的扩展是允许“partition1”和“partition2”由 jinja 模板变量定义，该变量包含从早期任务中提取的 XCOM，该任务将日期转换为财务年度和期间。使用日期作为分区是可能的，但我感兴趣的是是否可以强制 params 接受 Jinja 模板。

我想使用的代码如下所示：

        db = 'sample_db'
        table = 'sample_table'
        out = 's3://sample'
        p1 = '{{ task_instance.xcom_pull(task_ids="convert_to_partition", key="p1") }}'
        p2 = '{{ task_instance.xcom_pull(task_ids="convert_to_partition", key="p2") }}'

        add_partition_task= AWSAthenaOperator(
            task_id='add_partition',
            query='add_partition.sql',
            params={'database': db,
                    'table_name': table,
                    'p1': p1
                    'p2': p2},
            database=db,
            output_location=out
        )

所以现在 params.p1 和 params.p2 包含一个 Jinja 模板。显然，params 不支持 jinja 模板，因为渲染的 SQL 包含字符串文字 '{{ task_instance....' 而不是渲染的 XCOM 值。

将参数添加到运算符实现中的 template_fields 不足以强制其呈现模板。我的运算符仅扩展 BaseOperator 并使用扩展 AwsHook 的 AthenaHook。有谁有在参数（如结构）或替代方法中传递模板变量的经验吗？

Answer 1

由于AWSAthenaOperator既具有query模板化字段又接受文件扩展名.sql，因此您可以在文件本身中包含 jinja 模板。

我对您的 AWSAthenaOperator 进行了一些修改以适应该示例。

add_partition_task= AWSAthenaOperator(
    task_id='add_partition',
    query='add_partition.sql',
    params={
        'database': db,
        'table_name': table,
    }
)

这是add_partition.sql可能的样子。

INSERT OVERWRITE TABLE {{ params.database }}.{{ params.table_name }} (day, month, year) 
SELECT * FROM db.table
WHERE p1 = "{{ task_instance.xcom_pull(task_ids='convert_to_partition', key='p1') }}" 
  AND p2 = "{{ task_instance.xcom_pull(task_ids='convert_to_partition', key='p2') }}"
;