如何在Airflow DAG中通过SnowflakeOperator执行Snowpark Session代码？

SnowflakeOperator仅支持SQL脚本执行，要运行基于snowflake.Session的Snowpark DataFrame代码，最简便的方式是使用PythonOperator，同时复用Airflow中已配置的Snowflake连接信息，无需重复编写认证参数。

具体实现步骤 #

1. 复用Airflow Snowflake连接
   通过Airflow的BaseHook直接提取已配置的Snowflake连接信息，用来初始化snowflake.Session，避免硬编码账号、密码等敏感信息：

   from airflow import DAG
   from airflow.operators.python import PythonOperator
   from airflow.hooks.base import BaseHook
   import snowflake.snowpark as snowpark
   from snowflake.snowpark.functions import col
   from datetime import datetime
   
   def run_snowpark_dataframe_logic():
       # 替换为你在Airflow中配置的Snowflake连接ID
       conn_id = "your_snowflake_conn_id"
       snowflake_conn = BaseHook.get_connection(conn_id)
       
       # 构建Snowpark Session配置
       session_config = {
           "account": snowflake_conn.extra_dejson["account"],
           "user": snowflake_conn.login,
           "password": snowflake_conn.password,
           "warehouse": snowflake_conn.extra_dejson.get("warehouse"),
           "database": snowflake_conn.database,
           "schema": snowflake_conn.schema,
           "role": snowflake_conn.extra_dejson.get("role")
       }
       
       # 初始化Session并执行DataFrame操作
       with snowpark.Session.builder.configs(session_config).create() as session:
           # 示例：读取表、过滤数据并写入新表
           df = session.table("source_table").select(col("ID"), col("VALUE")).filter(col("VALUE") > 50)
           df.write.save_as_table("target_table", mode="append")
   
   # 现有DAG定义
   with DAG(
       dag_id="your_existing_dag",
       schedule_interval="@daily",
       start_date=datetime(2024, 1, 1),
       catchup=False
   ) as dag:
       # 原有的SnowflakeOperator任务
       existing_sql_task = SnowflakeOperator(
           task_id="execute_myscript_sql",
           sql="myscript.sql",
           snowflake_conn_id="your_snowflake_conn_id"
       )
       
       # 新增的Snowpark DataFrame任务
       snowpark_task = PythonOperator(
           task_id="run_snowpark_dataframe",
           python_callable=run_snowpark_dataframe_logic
       )
       
       # 设置任务依赖（根据实际需求调整）
       existing_sql_task >> snowpark_task
   

2. 环境依赖配置
   在Airflow运行环境中安装Snowpark Python包：

   pip install snowflake-snowpark-python
   

   确保包版本与你的Snowflake集群版本兼容。

3. 额外优化建议

   • 将Snowpark Session的初始化逻辑封装成独立工具函数，方便在多个任务中复用。
   • 通过PythonOperator的op_kwargs参数传递动态参数（如日期、表名）到DataFrame逻辑中。

关于SnowflakeOperator与Snowpark生态的适配 #

SnowflakeOperator的定位是批量执行SQL脚本，属于声明式的SQL任务调度；而Snowpark DataFrame API是Snowflake的编程式数据处理方案，两者面向不同的使用场景。目前Airflow官方没有提供专门的SnowparkOperator，因此使用PythonOperator结合Snowpark SDK是最直接的适配方式，既能利用Airflow的调度能力，又能完整使用Snowpark的生态功能。

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