Airflow：如何将Operator数据传入Snowflake存储过程（同会话访问临时表）

问题背景 #

需将Airflow Operator生成的批量数据传入SQL文件，创建临时表并确保后续存储过程能在同一会话访问该表，尝试三种方法均失败，具体代码及报错如下：

Operator核心代码 #

def list_of_tuples(self, worksheet):
    list_of_dicts = worksheet.get_all_records()
    reviewed_rows = [row for row in list_of_dicts]

    manual_data = [
        (
            str(row["NAME"]),
            str(row["VERSION"]),
            row.get("RSCORE", None),
            row["SERVERDATE"],
        )
        for row in reviewed_rows
    ]

    return manual_data

def run_proc(self, manual_data):
    source_hook = SnowflakeHook(snowflake_conn_id=self.source_conn_id)
   
    with closing(source_hook.get_conn()) as source_conn:
        with closing(source_conn.cursor(DictCursor)) as cur:
            try:
                cur.execute(self.sql, manual_data)

错误方法及原因 #

方法1：直接用Snowflake参数绑定批量数据

SQL文件：

EXECUTE IMMEDIATE $$
BEGIN
    CREATE TEMPORARY TABLE temp_manual_map (
        name STRING,
        version STRING,
        score NUMBER,
        serverdate DATE
    );

    INSERT INTO temp_manual_map
    VALUES :manual_map;
END;
$$;

报错：

TypeError: not all arguments converted during string formatting

原因：Snowflake Python Connector不支持将整个tuple列表作为单个参数传入VALUES子句，且参数传递逻辑与Python字符串格式化冲突，导致解析失败。

方法2：错误混用Jinja2与Snowflake语法

SQL文件：

EXECUTE IMMEDIATE $$
BEGIN
    CREATE TEMPORARY TABLE temp_manual_map (
        name STRING,
        version STRING,
        score INT,
        serverdate DATE
    );

    INSERT INTO manual_map (
        name,
        version,
        score,
        serverdate
    )
    VALUES
    {% for value in :manual_map %}
        (
            '{{ value[0] }}',
            '{{ value[1] }}',
            {{ value[2] }},
            '{{ value[3] }}'
        ){% if not loop.last %},{% endif %}
    {% endfor %};

END;
$$;

报错：

jinja2.exceptions.TemplateSyntaxError: unexpected ':'

原因：Jinja2模板引用变量无需加冒号，:manual_map是Snowflake参数语法，与Jinja2语法冲突。

方法3：无效SQL语句

SQL文件：

EXECUTE IMMEDIATE $$
BEGIN
    manual_map;
END;
$$;

报错：同方法1的字符串格式化错误
原因：manual_map不是合法SQL语句，参数传递逻辑仍未修正。

正确解决方案 #

方案1：使用Snowflake批量参数绑定（推荐，防SQL注入） #

修改Operator的run_proc方法，拆分批量数据为单个参数，用executemany实现批量插入，确保会话一致性：

def run_proc(self, manual_data):
    source_hook = SnowflakeHook(snowflake_conn_id=self.source_conn_id)
   
    with closing(source_hook.get_conn()) as source_conn:
        with closing(source_conn.cursor(DictCursor)) as cur:
            try:
                # 创建临时表
                cur.execute("""
                    CREATE TEMPORARY TABLE temp_manual_map (
                        name STRING,
                        version STRING,
                        score NUMBER,
                        serverdate DATE
                    );
                """)
                # 批量插入数据，使用%s作为占位符
                insert_sql = """
                    INSERT INTO temp_manual_map (name, version, score, serverdate)
                    VALUES (%s, %s, %s, %s);
                """
                # 执行批量插入
                cur.executemany(insert_sql, manual_data)
                # 同一会话内调用存储过程，可访问临时表
                # cur.execute("CALL your_target_procedure();")
            except Exception as e:
                source_conn.rollback()
                raise e
            else:
                source_conn.commit()

说明：全程在同一个数据库会话中执行，临时表可被后续操作直接访问，且批量插入方式避免SQL注入风险。

方案2：正确使用Jinja2模板渲染SQL（适配模板文件场景） #

修改SQL模板文件（移除冒号，添加SQL注入防护）

CREATE TEMPORARY TABLE temp_manual_map (
    name STRING,
    version STRING,
    score NUMBER,
    serverdate DATE
);

INSERT INTO temp_manual_map (name, version, score, serverdate)
VALUES
{% for row in manual_data %}
    (
        '{{ row[0] | replace("'", "''") }}',  -- 转义单引号防止注入
        '{{ row[1] | replace("'", "''") }}',
        {{ row[2] if row[2] is not none else 'NULL' }},
        '{{ row[3] }}'
    ){% if not loop.last %},{% endif %}
{% endfor %};

-- 同一会话内调用存储过程
-- CALL your_target_procedure();

修改Operator代码（渲染模板并执行）

from jinja2 import Template

def run_proc(self, manual_data):
    source_hook = SnowflakeHook(snowflake_conn_id=self.source_conn_id)
   
    # 读取并渲染SQL模板
    with open(self.sql_path, 'r') as f:
        sql_template = Template(f.read())
    rendered_sql = sql_template.render(manual_data=manual_data)
   
    with closing(source_hook.get_conn()) as source_conn:
        with closing(source_conn.cursor(DictCursor)) as cur:
            try:
                cur.execute(rendered_sql)
                source_conn.commit()
            except Exception as e:
                source_conn.rollback()
                raise e

说明：通过Jinja2渲染生成合法SQL，全程在同一会话执行，临时表可被后续存储过程访问。

内容的提问来源于stack exchange，提问作者Kar