1_opics_jtrust.py

import json, os
import csv
from datetime import date
import calendar

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator, ShortCircuitOperator
from airflow.operators.bash import BashOperator

from datetime import datetime, timedelta
from airflow.providers.sftp.operators.sftp import SFTPOperator
from airflow.models import Variable
from airflow.providers.postgres.operators.postgres import PostgresOperator

from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.utils.trigger_rule import TriggerRule

from airflow.models.xcom import XCom
from airflow.decorators import task
from airflow.hooks.postgres_hook import PostgresHook
from airflow import XComArg



# yesterday_nodash = (datetime.now() - timedelta(1)).strftime('%Y%m%d')
yesterday_nodash = (datetime.now() - timedelta(1)).strftime('%Y%m%d') if Variable.get(
    "DATE_OF_DATA") == 'today' else Variable.get("DATE_OF_DATA")
yesterday_strip = datetime.strptime(yesterday_nodash, '%Y%m%d').strftime('%Y-%m-%d')
d = f"""{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"""
today = d[:]

POSTGRES_CONN_ID = Variable.get("DS_DB")
POSTGRES_ENV = Variable.get("ENV_T24")
POSTGRES_SCHEMA = 'ds_opics' + Variable.get("ENV_T24")
# POSTGRES_SCHEMA = 'ds_switching'
DS_FOLDER = 'opics'
DS_DB = 'opics'
DS_SCHEMA = 'opics'
DS_CREATE_TABLE = ''


def _start():
    print("Start :: Extractor ")


def check_done_ext(**kwargs):
    from random import randint

    # number = randint(0, 10)
    print(f'''{Variable.get("LOCAL_PATH")}t24/{yesterday_nodash}/done.csv''')

    if os.path.isfile(f'''{Variable.get("LOCAL_PATH")}{DS_FOLDER}/{yesterday_nodash}/done_ext.csv'''):
        print("A")
        return False
    else:
        print("B")
        # STOP DAG
        return True


def ds_list_extractor():
    sql_stmt = f"""select * from ds_conf.ds_extractor_list_extractor('opics', '{yesterday_strip}');"""
    pg_hook = PostgresHook(
        postgres_conn_id=POSTGRES_CONN_ID,
    )
    pg_conn = pg_hook.get_conn()
    cursor = pg_conn.cursor()
    cursor.execute(sql_stmt)
    files = cursor.fetchall()
    return files

def ds_push_syntax(ti):
    iris = ti.xcom_pull(task_ids=['ds_list_extractor'])
    if not iris:
        raise Exception('No data.')

    return [{"op_kwargs": {
        "copy_sql": f"""COPY {POSTGRES_SCHEMA}.{table['table_name']} from STDOUT delimiter '{table['delimiter']}' CSV HEADER quote E'"'""",
        "file_id": f"""/{table['file_id']}"""}}
        for
        table in json.loads(iris[0][0][0])]

def pg_ddl_syntax(ti):
    iris = ti.xcom_pull(task_ids=['ds_list_extractor'])
    arr = []
    for table in json.loads(iris[0][0][0]):
        with open(f"""{Variable.get("LOCAL_PATH")}{DS_FOLDER}/{table['file_id']}""") as csvFile:
            reader = csv.reader(csvFile, delimiter=f"""{table['delimiter']}""")
            field_names_list = reader.__next__()
            arr.append({"sql": f"""drop table if exists {POSTGRES_SCHEMA}.{table['table_name']} cascade; create table {POSTGRES_SCHEMA}.{table['table_name']} ({' text, '.join([w.replace('.', '_').replace(' ', '_').replace('-', '_').replace("'", '').replace("(",'').replace(")",'').replace('/','or').lower() for w in field_names_list])} text);"""})
    return arr


def ds_push_csv(ti, copy_sql, file_id):
    pg_hook = PostgresHook.get_hook(POSTGRES_CONN_ID)
    pg_hook.copy_expert(copy_sql, filename=f"""/opt/airflow/dags/DFE/{DS_FOLDER}/{file_id}""")

with DAG("APJ_1_opics_jtrust",
         start_date=datetime(2021, 1, 1),
         schedule_interval=None,
         catchup=False,
         concurrency=3) as dag:
    begin = PythonOperator(
        task_id=f"Begin",
        python_callable=_start,
        op_kwargs={

        }
    )

    check_done_ext = ShortCircuitOperator(
        task_id="check_done_ext",
        provide_context=True,
        python_callable=check_done_ext,
        op_kwargs={},
    )

    history_start = PostgresOperator(
        sql=f"""INSERT INTO ds_conf."Datasource_history"
        (source, start_time, finish_time, status, env, date_of_data) VALUES
        ('{POSTGRES_SCHEMA}'::varchar(255), '{today}'::timestamp with time zone, '{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}'::timestamp with time zone, 'ONPROCESS'::varchar(100), '', '{yesterday_nodash}');""",
        task_id="history_start",
        postgres_conn_id=POSTGRES_CONN_ID,
    )

    drop_schema = PostgresOperator(
        sql=f"""drop schema if exists {POSTGRES_SCHEMA} cascade;""",
        task_id="drop_schema",
        postgres_conn_id=POSTGRES_CONN_ID,
    )

    create_schema = PostgresOperator(
        sql=f"""create schema if not exists {POSTGRES_SCHEMA};""",
        task_id="create_schema",
        postgres_conn_id=POSTGRES_CONN_ID,
    )


    ds_list_extractor = PythonOperator(
        task_id='ds_list_extractor',
        python_callable=ds_list_extractor,
        do_xcom_push=True
    )


    # pg_drop_schema = PostgresOperator(
    #     sql= f"""DROP TABLE IF EXISTS {POSTGRES_SCHEMA} CASCADE;""",
    #     task_id="pg_drop_schema",
    #     postgres_conn_id=POSTGRES_CONN_ID,
    # )
    #
    # pg_create_schema = PostgresOperator(
    #     sql= f"""create schema IF NOT EXISTS {POSTGRES_SCHEMA};""",
    #     task_id="pg_create_schema",
    #     postgres_conn_id=POSTGRES_CONN_ID,
    # )

    pg_ddl_syntax = PythonOperator(
        task_id='pg_ddl_syntax',
        python_callable=pg_ddl_syntax
    )

    pg_create_table = PostgresOperator.partial(
        task_id="pg_create_table",
        postgres_conn_id=POSTGRES_CONN_ID,
    ).expand_kwargs(
        XComArg(pg_ddl_syntax),
    )

    # ds_truncate_syntax = PythonOperator(
    #     task_id='ds_truncate_syntax',
    #     python_callable=ds_truncate_syntax
    # )
    #
    # ds_truncate = PostgresOperator.partial(
    #     task_id="ds_truncate",
    #     postgres_conn_id=POSTGRES_CONN_ID,
    # ).expand_kwargs(
    #     XComArg(ds_truncate_syntax)
    # )
    #
    ds_push_syntax = PythonOperator(
        task_id='ds_syntax_push',
        python_callable=ds_push_syntax
    )

    ds_csv_to_table = PythonOperator.partial(
        task_id="ds_csv_to_table",
        python_callable=ds_push_csv,
        dag=dag
    ).expand_kwargs(
        XComArg(ds_push_syntax),
    )

    set_access_schemma = PostgresOperator(
        sql=f"""GRANT USAGE ON SCHEMA {POSTGRES_SCHEMA} TO readaccess;""",
        task_id="set_access_schemma",
        postgres_conn_id=POSTGRES_CONN_ID,
    )

    set_access_all_table = PostgresOperator(
        sql=f"""GRANT SELECT ON ALL TABLES IN SCHEMA {POSTGRES_SCHEMA} TO readaccess;""",
        task_id="set_access_all_table",
        postgres_conn_id=POSTGRES_CONN_ID,
    )


    ds_ext_done = BashOperator(
        task_id="ds_ext_done",
        bash_command=f"""
          touch  /opt/airflow/dags/DFE/{DS_FOLDER}/{yesterday_nodash}/done_ext.csv;
          """,
    )

    history_finish = PostgresOperator(
        sql=f"""
               UPDATE ds_conf."Datasource_history"
                   SET status = 'DONE', finish_time = '{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}'::timestamp with time zone
               WHERE source = '{POSTGRES_SCHEMA}' and status = 'ONPROCESS';
    """,
        task_id="history_finish",
        postgres_conn_id=POSTGRES_CONN_ID,
    )




    begin >> check_done_ext >> history_start >> drop_schema >> create_schema >> ds_list_extractor >> pg_ddl_syntax >> pg_create_table >> ds_push_syntax >> ds_csv_to_table >> set_access_schemma >> set_access_all_table >> ds_ext_done >> history_finish