doczyai-pipelines/archive/airflow/dags/cicd_test_dag.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.


# Mandatory imports
from __future__ import annotations
import os
from datetime import datetime
import logging
from airflow import DAG
from airflow.operators.empty import EmptyOperator
from airflow.utils.trigger_rule import TriggerRule
from airflow.models import Variable

# SNOWFLAKE only imports
from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator

# PYTHON only imports
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.operators.python import PythonOperator


# Configuring basic logging for INFO / ERROR in our scripts

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Recommended to use this in the beginning of the script

SNOWFLAKE_CONN_ID = "doczy_dev_snowflake"  # Specific for every database
DAG_ID = "cicd_testing_dag"  # Must be unique for every dag
DATABASE = "XXX"

default_params = {"Database": "", "Schema": ""}

"""
Tags will be used to filter DAGs on the Airflow UI. 

Guidelines for TAGs:
"prod" - for dataload scripts that have been tested
"dev" - for dataloads for DEV databases / scripts in dev
"db_name" - same naming convention as Snowflake databses for clients / projects
"medical" / "pharmacy" - based on the scenario
"etl" / "adhoc" - based on the nature of the dataload

"""
TAGS = ["adhoc", "dev"]  # MANDATORY


"""
TRIGGER RULES for a task:

all_success: (default) all parents have succeeded
all_failed: all parents are in a failed or upstream_failed state
all_done: all parents are done with their execution
one_failed: fires as soon as at least one parent has failed, it does not wait for all parents to be done
one_success: fires as soon as at least one parent succeeds, it does not wait for all parents to be done
none_failed: all parents have not failed (failed or upstream_failed) i.e. all parents have succeeded or been skipped
none_skipped: no parent is in a skipped state, i.e. all parents are in a success, failed, or upstream_failed state
dummy: dependencies are just for show, trigger at will
"""
ALL_SUCCESS = "all_success"
ALL_FAILED = "all_failed"
ALL_DONE = "all_done"
ONE_SUCCESS = "one_success"
ONE_FAILED = "one_failed"


dag = DAG(
    # These args will get passed on to each operator
    # You can override them on a per-task basis during operator initialization
    # 'queue': 'bash_queue',
    # 'pool': 'backfill',
    # 'priority_weight': 10,
    # 'end_date': datetime(2016, 1, 1),
    # 'wait_for_downstream': False,
    # 'sla': timedelta(hours=2),
    # 'execution_timeout': timedelta(seconds=300),
    # 'on_failure_callback': some_function,
    # 'on_success_callback': some_other_function,
    # 'on_retry_callback': another_function,
    # 'sla_miss_callback': yet_another_function,
    DAG_ID,  # Mandatory for every dag
    start_date=datetime(2022, 1, 1),  # Must be in the past
    # Can pass snowflake conn id here instead of passing it to every task
    # It is needed when using Snowflake operator
    default_args={"snowflake_conn_id": SNOWFLAKE_CONN_ID, "retries": 0},
    tags=TAGS,
    catchup=False,  # True will run the dag on the specified frequency for backdated DAGs. Not applicate in out workflow
    schedule=None,  # If there is no fixed schedule, then always pass None explicitly
    params=default_params,
)


def python_op_eg(table_name, params):

    # Snowflake hook is used to fetch connection and cursor
    dwh_hook = SnowflakeHook(snowflake_conn_id=SNOWFLAKE_CONN_ID)
    conn = dwh_hook.get_conn()
    curr = conn.cursor()
    DATABASE = params["Database"]
    schema_name = params["Schema"]
    query_output = curr.execute(f"SELECT * from {DATABASE}.{schema_name}.{table_name}")

    # Alternative way to fetch query result

    # result = dwh_hook.get_first(f"select max(audit_sid) from {DATABASE}.stg.CLAIM_MED_STAGING")
    # max_audit_sid = result[0]

    logging.info(f"PYTHON OPERATOR OUTPUT = {query_output.fetchall()}")
    conn.close


# Best practice to have an empty start at the beginning and end
begin_job = EmptyOperator(task_id="Begin")

python_task = PythonOperator(
    task_id="get_info_using_python_op",  # Task ID has to be unique only inside a dags
    python_callable=python_op_eg,
    dag=dag,
    op_kwargs={
        "table_name": "DIM_AUDIT",
    },  # Variables can be passed to the python function using op_kwargs
)

# Snowflake operator does not offer the option to display query results
# We can use this for executing stored procedures.
sf_task = SnowflakeOperator(
    task_id="get_info_using_sf_op",
    sql=f"select * from {DATABASE}.STG.DIM_AUDIT",
    dag=dag,
)


end_job = EmptyOperator(task_id="End")

# EXECUTING TASKS
begin_job >> python_task >> sf_task >> end_job