Umang Mistry
156 lines
5.3 KiB
Python
156 lines
5.3 KiB
Python
#
|
|
# 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 |