1_los.py

import json, os
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') 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')
yesterday_lusa = (datetime.strptime(yesterday_nodash, '%Y%m%d') - timedelta(1)).strftime('%Y%m%d')
d = f"""{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"""
cron_tab = None if Variable.get("CRON_JOB") == 'manual' else Variable.get("CRON_JOB")

today = d[:]

POSTGRES_CONN_ID = Variable.get("DS_DB")
POSTGRES_SCHEMA = 'ds_los' + Variable.get("ENV_T24")
DS_CONN_ID = Variable.get("DS_LOS")
DS_FOLDER = 'los'
DS_DB = 'los'
DS_SCHEMA = 'public'
DS_CREATE_TABLE = ''


def create_ddl(schema, column):
    return f"""select string_agg(x,'') from (select
    'create table {POSTGRES_SCHEMA}."'||table_catalog||'_'||table_name|| '" ( ' || string_agg('"'||column_name||'"', ' text,' order by ordinal_position) || ' text'|| ' ); ' as x

from information_schema.columns where table_schema = '{schema}'
and table_name in ({column})
group by table_catalog, table_schema, table_name) as x ;"""


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


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

    # number = randint(0, 10)

    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('los', '{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_get_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 "{table['table_source']}".{table['file_id'].replace(".", "_")} to STDOUT delimiter '|' CSV header""",
        "file_id": f"""{table['table_source']}_{table['file_id']}""",
        "database": f"""{table['table_source']}"""}
        }
        for
        table in json.loads(iris[0][0][0])]


def ds_get_csv(ti, copy_sql, file_id, database):
    pg_hook = PostgresHook(
        postgres_conn_id=DS_CONN_ID,
        schema=database
    )
    print('INFO::',pg_hook)
    pg_hook.copy_expert(copy_sql, filename=f"""{Variable.get("LOCAL_PATH")}{DS_FOLDER}/""" + (
        f"""{yesterday_nodash}/{yesterday_nodash}.""" if Variable.get(
            "WITH_DATE") == 'Y' else '') + f"""{file_id}.csv""")


#
# def ds_drop_syntax(ti):
#     iris = ti.xcom_pull(task_ids=['ds_list_extractor'])
#     if not iris:
#         raise Exception('No data.')
#     return [{
#         "sql": f"""drop table ds_los.{table['file_id'].replace(".", "_")} CASCADE;"""}
#         for
#         table in json.loads(iris[0][0][0])]


def ds_get_ddl(ti):
    global DS_CREATE_TABLE
    iris = ti.xcom_pull(task_ids=['ds_list_extractor'])
    sql_table = []
    for table in json.loads(iris[0][0][0]):
        print(create_ddl(table['table_source'], f"""'{table['file_id']}'"""))
        sql_stmt = create_ddl(table['table_source'], f"""'{table['file_id']}'""")
        pg_hook = PostgresHook(
            postgres_conn_id=DS_CONN_ID,
            schema=table['table_source']
        )
        pg_conn = pg_hook.get_conn()
        cursor = pg_conn.cursor()
        cursor.execute(sql_stmt)
        sql_table.append(cursor.fetchall()[0][0])
    return sql_table


def pg_ddl_syntax(ti):
    print(ti.xcom_pull(task_ids=['ds_get_ddl'])[0])
    return [{"sql": ti.xcom_pull(task_ids=['ds_get_ddl'])[0]}]


def pg_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_source']}_{table['file_id'].replace(".", "_")}" from STDOUT delimiter '|' CSV header""",
        "file_id": f"""{table['table_source']}_{table['file_id']}""",}}
        for
        table in json.loads(iris[0][0][0])]


def pg_push_csv(ti, copy_sql, file_id):
    pg_hook = PostgresHook.get_hook(POSTGRES_CONN_ID)
    pg_hook.copy_expert(copy_sql, filename=f"""{Variable.get("LOCAL_PATH")}{DS_FOLDER}/""" + (
        f"""{yesterday_nodash}/{yesterday_nodash}.""" if Variable.get(
            "WITH_DATE") == 'Y' else '') + f"""{file_id}.csv""")


with DAG("APJ_1_los",
         start_date=datetime(2021, 1, 1),
         schedule=cron_tab,
         catchup=False,
         concurrency=3) as dag:
    begin = PythonOperator(
        task_id=f"Begin",
        python_callable=_start,
        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,
    )

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

    create_folder = BashOperator(
        task_id='create_folder',
        bash_command=f"""mkdir -p {Variable.get("LOCAL_PATH")}{DS_FOLDER}/{yesterday_nodash}"""
    )

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

    ds_get_syntax = PythonOperator(
        task_id='ds_syntax_get',
        python_callable=ds_get_syntax
    )

    ds_table_to_csv = PythonOperator.partial(
        task_id="ds_table_to_csv",
        python_callable=ds_get_csv,
        dag=dag
    ).expand_kwargs(
        XComArg(ds_get_syntax),
    )

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

    # ds_drop_syntax = PythonOperator(
    #     task_id='ds_drop_syntax',
    #     python_callable=ds_drop_syntax
    # )

    pg_drop_schema = PostgresOperator(
        sql=f"""drop schema 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),
    )

    pg_push_syntax = PythonOperator(
        task_id='pg_syntax_push',
        python_callable=pg_push_syntax
    )

    pg_csv_to_table = PythonOperator.partial(
        task_id="pg_csv_to_table",
        python_callable=pg_push_csv,
        dag=dag
    ).expand_kwargs(
        XComArg(pg_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;
         """,
    )

        
    ds_to_history = PostgresOperator(
        sql=f"""select ds_conf.ds_t24_create_table_history_surrounding('{yesterday_strip}', '{POSTGRES_SCHEMA}');""",
        task_id="ds_to_history",
        postgres_conn_id=POSTGRES_CONN_ID,
    )

    delete_before = BashOperator(
        task_id="delete_before",
        bash_command=f"""
         rm -r /opt/airflow/dags/DFE/{DS_FOLDER}/{yesterday_lusa};
         """,
    )

    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 >> create_folder >> ds_list_extractor >> ds_get_syntax >> ds_table_to_csv >> ds_get_ddl >> pg_drop_schema >> pg_create_schema >> pg_ddl_syntax >> pg_create_table >> pg_push_syntax >> pg_csv_to_table >> set_access_schemma >> set_access_all_table >> ds_ext_done  >> ds_to_history >> delete_before >> history_finish