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 }}')

此模板效果很好。

对此的扩展是允许“ jina”模板变量定义“ partition1”和“ partition2”，该变量包含来自较早任务的XCOM拉取，该任务将日期转换为“财政年度和期间”。可以使用date作为分区，但是我对是否可以强制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') }}"
;

气流Jinja在参数中进行模板化

1 个答案: