Notes

import os
from datetime import datetime
from typing import Dict
import logging
import sys
from airflow import DAG
from airflow.decorators import task
from airflow.operators.python import BranchPythonOperator
from airflow.operators import bash_operator
from airflow.providers.google.cloud.operators.gcs import GCSListObjectsOperator
from airflow.providers.google.cloud.sensors.gcs import GCSObjectsWithPrefixExistenceSensor, GCSObjectExistenceSensor
from airflow.contrib.operators.gcs_to_bq import GoogleCloudStorageToBigQueryOperator
from airflow.providers.google.cloud.transfers.gcs_to_gcs import GCSToGCSOperator
from hrastro.plugins.operators.CreateDataProcCluster import CreateDataProcCluster
from hrastro.plugins.operators.DataprocSparkOperator import DataprocSparkOperator
from hrastro.plugins.operators.DeleteDataProcCluster import DeleteDataProcCluster
from hrastro.plugins.utils import CommonUtil
from hrastro.plugins.utils import appConstants as AC
from hrastro.plugins.utils import stringConstants as SC
from hrastro.plugins.operators import DataprocBashOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from airflow.operators.dummy import DummyOperator
from airflow.operators.bash import BashOperator

DEFAULT_ARG: Dict = {
SC.GCP_CONN_ID: 'gcp_etl_conn_id'
, 'depends_on_past': False
, SC.PROJECT_ID: CreateDataProcCluster.get_project_id_from_connection(gcp_conn_id='gcp_etl_conn_id')
, SC.REGION: CreateDataProcCluster.get_dataproc_region()
, SC.DOMAIN_NAME: AC.WS_VAR.get(SC.DOMAIN_NAME)
, SC.OWNER: AC.WS_VAR.get(SC.OWNER)
, SC.RUN_BOOK: ''
, SC.TR_PRODUCT_ID: AC.WS_VAR.get(SC.TR_PRODUCT_ID)
, SC.DAG_FILE_NAME: os.path.splitext(os.path.basename(__file__))[0]
, SC.APPLICATION: 'example'
, SC.PRIORITY: 'P1'
, SC.PRODUCT_LINE: "example"
}

#
dag_id: str = CommonUtil.get_dag_id(DEFAULT_ARG)
DEFAULT_ARG[SC.CLUSTER_NAME] = CommonUtil.get_cluster_name(dag_id)

dag: DAG = DAG(
dag_id=dag_id
, catchup=False
, start_date=datetime(2021, 5, 11)
, schedule_interval=None
, max_active_tasks=20
, max_active_runs=20
, default_args=DEFAULT_ARG
)


# [START instantiate_dag]

@task()
def get_report_name(data: list):
report_dict = {}
report_full_path: str = data.pop()
report_full_name: str = report_full_path.split("/").pop()
report_type: str = report_full_name.split(".")[1].lower()
report_name: str = report_full_name.split(".")[0]
report_split = report_name.split("_")

if(len(report_split)==3):
gcp_report_id: str = report_split[0].lower()
report_variant: str = report_split[1]
report_state_code: str = ""
report_transaction_ts: str = report_split[2]
elif(len(report_split)==4):
gcp_report_id: str = report_split[0].lower()
report_variant: str = report_split[1]
report_state_code: str = report_split[2]
report_transaction_ts: str = report_split[3]

default_report_id: str = gcp_report_id[3:]
report_id: int = int(default_report_id.lstrip("0"))
transaction_date: str = report_transaction_ts[0:8]
report_transaction_date = datetime.strptime(transaction_date, "%Y%m%d").strftime('%Y-%m-%d')
report_transaction_timestamp = datetime.strptime(report_transaction_ts, '%Y%m%d%H%M%S').strftime('%Y-%m-%d %H:%M:%S')
raw_table_name = "sap_report_data_"+ gcp_report_id
base_table_name = "sap_report_data_"+str(report_id)+"_"+transaction_date[0:4]
gcs_report_full_path = "gs://"+AC.BUCKET_VAR.get("landing_payroll_bucket")+"/"+report_full_path
report_dict['source_file_full_path'] = report_full_path
report_dict['gcs_report_full_path'] = gcs_report_full_path
report_dict['source_path'] = "gs://{source_bucket}/{report_full_path}".format(source_bucket=AC.BUCKET_VAR.get("landing_payroll_bucket"),report_full_path=report_full_path)
report_dict['report_id'] = report_id
report_dict['report_variant'] = report_variant
report_dict['gcp_report_id'] = gcp_report_id
report_dict['report_state_code'] = report_state_code
report_dict['report_transaction_ts'] = report_transaction_ts
report_dict['report_transaction_timestamp'] = report_transaction_timestamp
report_dict['transaction_date'] = transaction_date
report_dict['base_transaction_date'] = report_transaction_date
report_dict['report_type'] = report_type
report_dict['report_name'] = report_name
report_dict['raw_table_name'] = raw_table_name
report_dict['base_table_name'] = base_table_name
report_dict['report_full_name'] = report_full_name
report_dict['bq_table_name'] = base_table_name.upper()
if(report_type == 'txt' and report_id == 4):
report_dict['jar_file_name'] = "PayrollCustomWorkflow04"
elif(report_type=='txt'):
report_dict['jar_file_name'] = "PayrollTextFilesWorkflow"
elif(report_type=='csv' or report_type == 'orc'):
report_dict['jar_file_name'] = "CsvPayrollWorkflow"

report_dict['dest_path'] = "{db_name}.db/{table_name}/data_date={transaction_date}/{file_name}".format(
db_name="ww_hr_dl_raw_payroll"
, table_name=raw_table_name
, transaction_date=transaction_date
, file_name=report_full_name)

download_link = "https://storage.cloud.google.com/{raw_bucket}/{rawSchema}.db/{raw_table_name}/data_date={transaction_date}/{file_name}".format(
raw_bucket=AC.BUCKET_VAR.get('raw_payroll_bucket')
, rawSchema="ww_hr_dl_raw_payroll"
, raw_table_name=raw_table_name
, transaction_date=transaction_date
, file_name=report_full_name)
report_dict['download_link'] = download_link
return report_dict


def decide_load_type(**kwargs):
import json
data = kwargs['templates_dict']['data'].replace("'", """)
json_data = json.loads(data)
report_type = json_data.get("report_type")
if report_type in ['csv', 'txt', 'orc']:
return 'hr_create_cluster'
elif report_type in ['pdf']:
return 'hr_bq_data_load_task'
else:
return 'end_task'

check_load_type = BranchPythonOperator(
task_id='check_load_type',
python_callable=decide_load_type,
provide_context=True,
templates_dict={'data': "{{ ti.xcom_pull(task_ids='get_report_name', key='return_value')}}"},
dag=dag
)

bq_data_load_task = BigQueryInsertJobOperator(
task_id="hr_bq_data_load_task",
job_id="base_to_bq_query_id001",
configuration={
"query": {
"query": "INSERT INTO `wmt-us-hr-dev.WW_HR_DL_US_PAYROLL_REPORT_TABLES.{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('bq_table_name') }}` (report_id,search_key_1,report_variant,updated_on,data,transaction_date) VALUES ({{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_id') }},'{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_state_code')}}','{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_state_code')}}_{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_variant')}}','{{ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_transaction_timestamp')}}','{{ti.xcom_pull(task_ids='get_report_name', key='return_value').get('download_link')}}','{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('base_transaction_date')}}');",
"useLegacySql": False,
}
},
location="US",
dag=dag
)
gcsPath="{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('gcs_report_full_path') }}"
scripts_path="/datapipeline/application/payroll/script"
delimiter="t"
rownum_task = DataprocBashOperator.CustomDProcBashOperator(
task_id='row_num_generator',
bash_script='generate_rownum.sh',
dag=dag,
arguments=(gcsPath,delimiter),
jar_file_uris=['gs://99d3a35648706bb0f9775de0077c88b8e6805486e59dff4eeb234436e21b64/datapipeline/application/payroll/script/generate_rownum.sh']
)

bq_map_load_task = BigQueryInsertJobOperator(
task_id="bq_map_load_task",
job_id="base_to_bq_query_id002",
configuration={
"query": {
"query": "INSERT INTO `wmt-us-hr-dev.WW_HR_DL_US_PAYROLL_REPORT_TABLES.SAP_REPORT_FILTER` (report_id,search_keys,report_display_name,report_variant,updated_on,download_link,transaction_date,report_header) VALUES ({{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_id') }},STRUCT([STRUCT<key STRING, value STRING>('search_key_1' , 'State Code'),STRUCT<key STRING, value STRING>('search_key_2' , 'Not Applicable'),STRUCT<key STRING, value STRING>('search_key_3' , 'Not Applicable'),STRUCT<key STRING, value STRING>('search_key_4' , 'Not Applicable')]),'{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_state_code')}}_{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_variant')}}','{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_state_code')}}_{{ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_variant')}}','{{ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_transaction_timestamp')}}','{{ti.xcom_pull(task_ids='get_report_name', key='return_value').get('download_link')}}','{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('base_transaction_date')}}','Download Link');"
, "useLegacySql": False,
}
},
location="US",
dag=dag
)

gcs_object_with_prefix_exists = GCSObjectsWithPrefixExistenceSensor(
google_cloud_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID),
bucket=AC.BUCKET_VAR.get("landing_payroll_bucket"),
prefix="sap_payroll_report/us/sap/GCP",
mode='poke',
task_id="check_any_payroll_reports",
dag=dag
)

list_payroll_reports = GCSListObjectsOperator(
task_id="list_payroll_reports"
, bucket=AC.BUCKET_VAR.get("landing_payroll_bucket")
, prefix="sap_payroll_report/us/sap/GCP"
, dag=dag
)

single_object = get_report_name(list_payroll_reports.output)

gcs_object_exists = GCSObjectExistenceSensor(
google_cloud_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID),
bucket=AC.BUCKET_VAR.get("landing_payroll_bucket"),
object="{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('source_file_full_path') }}",
mode='poke',
task_id="check_payroll_report",
dag=dag
)
move_files_archive = GCSToGCSOperator(
task_id="move_files_archive",
source_bucket=AC.BUCKET_VAR.get("landing_payroll_bucket"),
source_object="{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('source_file_full_path') }}",
destination_bucket=AC.BUCKET_VAR.get('landing_payroll_bucket'),
destination_object="sap_payroll_report/us/sap/archive/{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_full_name') }}",
gcp_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID),
move_object=False,
dag=dag
)
move_single_file = GCSToGCSOperator(
task_id="move_files_copy",
source_bucket=AC.BUCKET_VAR.get("landing_payroll_bucket"),
source_object="{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('source_file_full_path') }}",
destination_bucket=AC.BUCKET_VAR.get('raw_payroll_bucket'),
destination_object="{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('dest_path') }}",
gcp_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID),
move_object=True,
dag=dag
)


create_cluster = CreateDataProcCluster(
task_id="create_cluster"
, dag=dag
)

delete_cluster = DeleteDataProcCluster(
dag=dag)

spark_task = DataprocSparkOperator(
jar_file_uris=["gs://{}/datapipeline/application/payroll/bin/payroll-2.0.9-SNAPSHOT.jar".format(AC.BUCKET_VAR.get(SC.SRC_BUCKET))
, "gs://2a84d15336bfbfecadcd8edc4efdf6f9371899f62e5b8ed7a73f0794bcf016/datapipeline/application/sparkjars/jars_com.jcraft_jsch-0.1.53.jar"
, "gs://2a84d15336bfbfecadcd8edc4efdf6f9371899f62e5b8ed7a73f0794bcf016/datapipeline/application/sparkjars/jars_io.delta_delta-core_2.11-0.5.0.jar"]
, main_class="com.walmart.dsi.hr.payroll.WorkflowController"
, arguments=["workflow={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('jar_file_name') }}"
, "runmode=global"
, "rawSchema=ww_hr_dl_raw_payroll"
, "baseSchema=ww_hr_dl_base_payroll"
, "report_id={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_id') }}"
, "transaction_date={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('transaction_date') }}"
, "rawBucketId={}".format(AC.BUCKET_VAR.get('raw_payroll_bucket'))
, "source_file_name={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_name') }}"
, "jsonFileName=PayrollMappingJson"
, "etlRunDateTime={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_transaction_ts') }}"
, "reportVariant={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('report_variant') }}"
, "gcpReportName={{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('gcp_report_id') }}"]
, dag=dag
)

base_to_bq_data_task = GoogleCloudStorageToBigQueryOperator(
task_id="base_to_bq_data_task"
, bucket= AC.BUCKET_VAR.get('base_payroll_bucket')
, source_objects=["ww_hr_dl_base_payroll.db/{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('base_table_name') }}/part*"]
, destination_project_dataset_table="WW_HR_DL_US_PAYROLL_REPORT_TABLES.{{ ti.xcom_pull(task_ids='get_report_name', key='return_value').get('bq_table_name') }}"
, source_format='PARQUET'
, schema_object = ""
, create_disposition='CREATE_IF_NEEDED'
, write_disposition='WRITE_APPEND'
, bigquery_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID)
, google_cloud_storage_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID)
, dag=dag
)

base_to_bq_mapping_task = GoogleCloudStorageToBigQueryOperator(
task_id="base_to_bq_mapping_task"
, bucket= AC.BUCKET_VAR.get('base_payroll_bucket')
, source_objects=["ww_hr_dl_base_payroll.db/sap_report_filter/part*"]
, destination_project_dataset_table="WW_HR_DL_US_PAYROLL_REPORT_TABLES.SAP_REPORT_FILTER"
, source_format='PARQUET'
, schema_object = ""
, create_disposition='CREATE_IF_NEEDED'
, write_disposition='WRITE_APPEND'
, bigquery_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID)
, google_cloud_storage_conn_id=DEFAULT_ARG.get(SC.GCP_CONN_ID)
, dag=dag
)


start_task = DummyOperator(task_id='start_task', dag=dag)
end_task = DummyOperator(task_id='end_task', dag=dag)

start_task >> gcs_object_with_prefix_exists >> list_payroll_reports
single_object >> gcs_object_exists

gcs_object_exists >> create_cluster >> rownum_task >> move_files_archive >> move_single_file >> check_load_type
check_load_type >> bq_data_load_task >> bq_map_load_task
check_load_type >> spark_task >> base_to_bq_data_task >> base_to_bq_mapping_task
[bq_map_load_task, base_to_bq_mapping_task] >> end_task