Apache Airflow Snowflake Provider

from __future__ import annotations
import os
import pendulum

from airflow.models.dag import DAG
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator

SNOWFLAKE_CONN_ID = os.environ.get('SNOWFLAKE_CONN_ID', 'snowflake_default')

with DAG(
    dag_id='snowflake_quickstart_example',
    start_date=pendulum.datetime(2023, 1, 1, tz='UTC'),
    catchup=False,
    schedule=None,
    tags=['snowflake', 'example'],
    doc_md="""### Snowflake Quickstart DAG
    This DAG demonstrates a basic interaction with Snowflake using the SnowflakeOperator.
    It creates a table, inserts data, and then queries the table.
    """,
) as dag:
    create_table = SnowflakeOperator(
        task_id='create_snowflake_table',
        snowflake_conn_id=SNOWFLAKE_CONN_ID,
        sql=(
            "CREATE TABLE IF NOT EXISTS AIRFLOW_TEST_TABLE (id INTEGER, name VARCHAR);"
        ),
    )

    insert_data = SnowflakeOperator(
        task_id='insert_data_into_table',
        snowflake_conn_id=SNOWFLAKE_CONN_ID,
        sql="INSERT INTO AIRFLOW_TEST_TABLE (id, name) VALUES (1, 'Airflow'), (2, 'Snowflake');",
    )

    query_data = SnowflakeOperator(
        task_id='query_snowflake_table',
        snowflake_conn_id=SNOWFLAKE_CONN_ID,
        sql="SELECT * FROM AIRFLOW_TEST_TABLE;",
        handler=lambda cursor: [print(row) for row in cursor.fetchall()],
    )

    # Define task dependencies
    create_table >> insert_data >> query_data