from datetime import datetime
from airflow import DAG

from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator
from airflow.providers.common.sql.operators.sql import SQLColumnCheckOperator

my_conn_id = "snowflake_default"

default_args={"conn_id": my_conn_id}

with DAG(
    dag_id="airflow_providers_example",
    schedule=None,
    start_date=datetime(2022, 11, 27),
    default_args=default_args,
) as dag:

    create_table = SnowflakeOperator(
        task_id="create_table",
        sql=""" CREATE OR REPLACE TABLE testing AS (
                        SELECT
                            1 AS row_num,
                            'not null' AS field

                        UNION ALL

                        SELECT
                            2 AS row_num,
                            'test' AS field

                        UNION ALL

                        SELECT
                            3 AS row_num,
                            'test 2' AS field
                    )""",
    )

    column_checks = SQLColumnCheckOperator(
        task_id="column_checks",
        table="testing",
        column_mapping={
            "field": {"unique_check": {"equal_to": 0}, "null_check": {"equal_to": 0}}
        },
    )

    create_table >> column_checks