Leverage the Power of Airflow To Build an End-to-End Data Pipeline

				
					"""
Variables are used in this script:
==================================
1. windsor_api_key
2. windsor_api
3. gcp_weselect_project
4. gcp_welect_bigquery_campaign_dataset


Connections are used in this script:
==================================
1. windsor_api_available
2. weselect_bigquery_gcp_conn
"""


import os
import json
import requests
import pandas as pd
import traceback
from datetime import date, timedelta, datetime
from google.oauth2 import service_account
from google.cloud import bigquery
from airflow.models import DAG
from airflow.utils.task_group import TaskGroup
from airflow.providers.http.sensors.http import HttpSensor
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryCreateEmptyDatasetOperator,\
   BigQueryCreateEmptyTableOperator
from airflow.operators.bash import BashOperator
from airflow.models import Variable
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.edgemodifier import Label


dag_bucket = Variable.get('dag_bucket')
weselect_data_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'weSelect')
tiktok_dir = f'gs://{dag_bucket}/dags/weSelect/campaign_data/tiktok/'
yesterday = date.today() - timedelta(days = 1) # for last day data capture from windsor API


default_args = {
   'owner': 'weSelect_admin',
   'email': [‘user_mail@example.com‘], ‘retries’: 5, ‘retry_delay’: timedelta(seconds=20) ],
   'retries': 5,
   'retry_delay': timedelta(seconds=120)
}


def _check_api_data(ti):
   """
   check windsor tiktok api data is available or not
   """
   api_endpoint = f"{Variable.get('windsor_api')}/tiktok?api_key={Variable.get('windsor_api_key')}&fields=ctr%2Caccount_id%2Caccount_name%2Cad_name%2Ccampaign%2Ccampaign_name%2Ccampaign_id%2Cclicks%2Cdate%2Cduet_clicks%2Cimpressions%2Cplay_first_quartile%2Creach%2Csound_usage_clicks%2Cspend%2Cstitch_clicks%2Cad_id&date_from={yesterday}&date_to={yesterday}"
   response = requests.get(api_endpoint)
   data = json.loads(response.content)['data']
   print(f"data volume: {len(data)}")
   if len(data) > 0:  # if data is available then export data as csv
       ti.xcom_push(key='campaign_data', value=data)
       return 'dump_campaign_data'
   else:  # otherwise, skip
       return 'empty_data'


def _dump_campaign_data(ti):
   """
   export campaign data as csv
   """


   campaign_data = ti.xcom_pull('check_api_data', key='campaign_data')  # pull API data
   df = pd.DataFrame(campaign_data)  # convert to pandas dataframe
   file_name = f'campaign_data_{datetime.now()}.csv'
   print(f"total campaign data: {len(campaign_data)}")
   print('Data is exported as CSV successfully.')
   df.to_csv(f'{tiktok_dir}{file_name}', index=False)


def upload_data_from_csv():
   """
   fetch data from csv and upload to bigQuery table
   """


   bgquery_client = bigquery.Client()  # bigquery


   # create a job
   job_config = bigquery.LoadJobConfig(source_format=bigquery.SourceFormat.CSV
                                       , skip_leading_rows=1
                                       , autodetect=False
                                       , write_disposition=bigquery.WriteDisposition.WRITE_APPEND
                                       )
   try:
       df = pd.read_csv(f'{tiktok_dir}*.csv')  # tiktok dir
       df = df.where(pd.notnull(df), None)  # convert pandas NaN to SQL null
       print(f"data: {df.head(5)}")


       job = bgquery_client.load_table_from_dataframe(df
                                       , f"{Variable.get('gcp_weselect_project')}.{Variable.get('gcp_welect_bigquery_campaign_dataset')}.tiktok_campaign"
                                       , job_config=job_config)
       job.result()  # Waits for the job to complete
       print(f"Data insertion is successful in tiktok_campaign table.")
   except Exception:
       print(f"Something is wrong with csv_file file and tiktok_campaign table.")
       print(f"Error: {traceback.print_exc()}")


with DAG(dag_id='weSelect_windsor_tiktok_campaign_data_composer'
        , tags=['gcp_bigquery', 'windsor', 'tiktok', 'campaign']
        , default_args=default_args
        , description=f'fetch all tiktok campaign data from Windsor API if available and save the {yesterday} data int GCP BigQuery'
        , start_date=datetime(2024, 2, 24)
        , schedule='0 5 * * *'  # 5 AM
        , catchup=False):
  
   # Tasks
   # =====
  
   #1. check API availability
   is_windsor_tiktok_api_available = HttpSensor(task_id='is_windsor_api_available'
                                                  , http_conn_id='windsor_api_available'
                                                  , endpoint=f'tiktok?api_key={Variable.get("windsor_api_key")}&fields=campaign&date_from={yesterday}&date_to={yesterday}'
                                                  , response_check=lambda response: True if response.status_code == 200 else False
                                                  , method='GET'
                                                  , mode='poke'
                                                  , poke_interval=5
                                                  , timeout=60
                                                  )
  
   #2. fetch campaign data
   check_api_data = BranchPythonOperator(task_id='check_api_data'
                                         , python_callable=_check_api_data
                                         , trigger_rule=TriggerRule.ALL_SUCCESS
                                         )
  
   #3. empty data
   empty_data = PythonOperator(task_id='empty_data'
                               , python_callable=lambda: print(f"No data is available in API for {yesterday}")
                               )
  
   #4. dump campaign data
   dump_campaign_data = PythonOperator(task_id='dump_campaign_data'
                                       , python_callable=_dump_campaign_data
                                       )
  
   #5. create BigQuery dataset and table
   with TaskGroup(group_id='create_bigquery_dataset_table'
                  , tooltip='create campaign dataset and tiktok_campaign table') as create_bigquery_dataset_table:
      
       #5.1 create dataset
       create_campaign_data_dataset = BigQueryCreateEmptyDatasetOperator(task_id='create_campaign_data_dataset',
                                                                       gcp_conn_id='weselect_bigquery_gcp_conn',
                                                                       project_id='{{ var.value.gcp_weselect_project}}',
                                                                       dataset_id='{{ var.value.gcp_welect_bigquery_campaign_dataset}}',
                                                                       if_exists='ignore'
                                                                       )
      
       #5.2 create table
       with open(f'{weselect_data_dir}/table_schema/tiktok_schema.json') as f:  # tiktok_campaign table schema
           schema = json.load(f)


       create_tiktok_campaign_table = BigQueryCreateEmptyTableOperator(task_id='create_tiktok_campaign_table',
                                                               gcp_conn_id='weselect_bigquery_gcp_conn',
                                                               project_id='{{ var.value.gcp_weselect_project}}',
                                                               dataset_id='{{ var.value.gcp_welect_bigquery_campaign_dataset}}',
                                                               table_id='tiktok_campaign',
                                                               if_exists='ignore',
                                                               schema_fields=schema,
                                                               time_partitioning={
                                                                   "type": "YEAR",
                                                                   "field": "date"
                                                                   }
                                                                   )
       # dataset and table taskgroup taskflow
       create_campaign_data_dataset >> create_tiktok_campaign_table
  
   # #6. insert data into BigQuery
   upload_data = PythonOperator(task_id='upload_data'
                                , python_callable=upload_data_from_csv)




   # Taskflow
   [create_bigquery_dataset_table]


   is_windsor_tiktok_api_available \
   >> Label('check windsor tiktok API is available') >> check_api_data
   check_api_data >> Label('data unavaiable') >> empty_data
   check_api_data >> Label('data available') >> dump_campaign_data \
   >> Label('export API data as csv') >> upload_data \
   >> Label('append data to BigQuery')