Таблица-справочник – генератор DAG? А что так можно было?

from datetime import datetime, timedelta

import pendulum

from airflow import DAG

from airflow.operators.empty import EmptyOperator
from airflow.operators.python import PythonOperator

from airflow.providers.postgres.hooks.postgres import PostgresHook

# Конфигурация DAG
OWNER = 'korsak0v'
DAG_ID = 'check_pg'
LOCAL_TZ = pendulum.timezone('Europe/Moscow')

# Описание возможных ключей для default_args
# https://github.com/apache/airflow/blob/343d38af380afad2b202838317a47a7b1687f14f/airflow/example_dags/tutorial.py#L39
args = {
    'owner': OWNER,
    'start_date': datetime(2023, 1, 1, tzinfo=LOCAL_TZ),
    'catchup': True,
    'retries': 3,
    'retry_delay': timedelta(hours=1),
}


def check_pg_connect(**context):
    """"""
    pg = PostgresHook('test_db')

    df = pg.get_pandas_df('SELECT 1 AS one')

    if len(df) == 1:
        print(True)

with DAG(
        dag_id=DAG_ID,
        schedule_interval='10 0 * * *',
        default_args=args,
        tags=['check_pg_connect', 'test'],
        concurrency=1,
        max_active_tasks=1,
        max_active_runs=1,
) as dag:

    start = EmptyOperator(
        task_id='start',
    )

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

    end = EmptyOperator(
        task_id='end',
    )

    start >> check_pg_connect >> end

from datetime import datetime, timedelta

import pendulum

from airflow import DAG

from airflow.operators.empty import EmptyOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator


# Конфигурация DAG
OWNER = 'korsak0v'
DAG_ID = 'simple_dag'
LOCAL_TZ = pendulum.timezone('Europe/Moscow')

# Названия коннекторов к PG
PG_CONNECT = 'test_db'

# Используемые таблицы в DAG
PG_TARGET_SCHEMA = 'dm'
PG_TARGET_TABLE = 'fct_sales'
PG_TMP_SCHEMA = 'stg'
PG_TMP_TABLE = f'tmp_{PG_TARGET_TABLE}_{{{{ data_interval_start.format("YYYY_MM_DD") }}}}'
INDEX_KPI = 1

sql_query = '''
SELECT 
	('2023-'||((random()*11+1)::int)::varchar||'-'||((random()*27+1)::int)::varchar)::date AS date,
	(random()*100)::int AS value,
	1 AS kpi_id
'''

LONG_DESCRIPTION = '# LONG_DESCRIPTION'

SHORT_DESCRIPTION = 'SHORT_DESCRIPTION'


args = {
    'owner': OWNER,
    'start_date': datetime(2023, 1, 1, tzinfo=LOCAL_TZ),
    'catchup': True,
    'retries': 3,
    'retry_delay': timedelta(hours=1),
}

with DAG(
        dag_id=DAG_ID,
        schedule_interval='10 0 * * *',
        default_args=args,
        tags=['dm'],
        description=SHORT_DESCRIPTION,
        concurrency=1,
        max_active_tasks=1,
        max_active_runs=1,
) as dag:
    dag.doc_md = LONG_DESCRIPTION

    start = EmptyOperator(
        task_id='start',
    )

    drop_tmp_before = PostgresOperator(
        task_id='drop_tmp_before',
        sql=f'''DROP TABLE IF EXISTS {PG_TMP_SCHEMA}.{PG_TMP_TABLE}''',
        postgres_conn_id=PG_CONNECT
    )

    create_tmp = PostgresOperator(
        task_id='create_tmp',
        sql=f'''
        CREATE TABLE {PG_TMP_SCHEMA}.{PG_TMP_TABLE} AS
        {
            sql_query.format(
                start_date="{{ data_interval_start.format('YYYY-MM-DD') }}",
                end_date="{{ data_interval_end.format('YYYY-MM-DD') }}"
            )
        };
        ''',
        postgres_conn_id=PG_CONNECT
    )

    delete_from_target = PostgresOperator(
        task_id='delete_from_target',
        sql=f'''
        DELETE FROM {PG_TARGET_SCHEMA}.{PG_TARGET_TABLE}
        WHERE 
            date IN (
                SELECT 
                    date 
                FROM 
                    {PG_TMP_SCHEMA}.{PG_TMP_TABLE}
                )
        ''',
        postgres_conn_id=PG_CONNECT
    )

    insert_from_tmp_to_target = PostgresOperator(
        task_id='insert_from_tmp_to_target',
        sql=f'''
        INSERT INTO {PG_TARGET_SCHEMA}.{PG_TARGET_TABLE}("date", value, kpi_id)
        SELECT 
            "date", 
            value,
            kpi_id
        FROM 
            {PG_TMP_SCHEMA}.{PG_TMP_TABLE}
        ''',
        postgres_conn_id=PG_CONNECT
    )

    drop_tmp_after = PostgresOperator(
        task_id='drop_tmp_after',
        sql=f'''DROP TABLE IF EXISTS {PG_TMP_SCHEMA}.{PG_TMP_TABLE}''',
        postgres_conn_id=PG_CONNECT
    )

    end = EmptyOperator(
        task_id='end',
    )

    start >> drop_tmp_before >> create_tmp >> delete_from_target >> insert_from_tmp_to_target >> drop_tmp_after >> end
    

CREATE SCHEMA dm;

CREATE SCHEMA stg;

DROP TABLE IF EXISTS dm.fct_sales;

CREATE TABLE dm.fct_sales (
	id bigserial PRIMARY KEY,
	"date" date NOT NULL,
	kpi_id int2 NOT NULL,
	value int4 NOT NULL
);

CREATE TABLE public.dim_kpi_dag_gen_config (
	id serial4 NOT NULL,
	kpi_id int4 NOT NULL,
	dag_id varchar(255) NOT NULL,
	"owner" varchar(30) NOT NULL,
	start_date varchar(40) NOT NULL,
	metric_name_en varchar(255) NOT NULL,
	sql_query text NOT NULL,
	short_description_md varchar(255) NULL,
	long_description_md text NULL,
	cron varchar(50) NOT NULL,
	sensors varchar(255) NULL,
	tags varchar(255) NOT NULL,
	metric_line varchar(255) NULL,
	"source" varchar(255) NULL,
	bi_logic varchar(255) NULL,
	comment_pa text NULL,
	comment_de text NULL,
	is_actual bool NULL DEFAULT TRUE,
	created_at timestamptz NULL DEFAULT now(),
	changed_at timestamptz NULL,
	pg_environment varchar(10) NOT NULL DEFAULT 'prod'::character varying,
	airflow_environment varchar(10) NOT NULL DEFAULT 'dev'::character varying,
	CONSTRAINT dim_kpi_dag_gen_config_pkey PRIMARY KEY (id)
)

CREATE TABLE public.fct_dm_kpi (
	"date" date NULL,
	value float8 NULL,
	kpi_id int4 NULL
);

from connectors_to_databases import PostgreSQL


pg = PostgreSQL(
    port=1
)


TABLE = 'dim_kpi_dag_gen_config'


def gen_insert_sql_for_kpi_id(dict_kpi: dict = None) -> str:
    """
    Генерирует скрипт для вставки данных в SCD.

    Определяет есть ли такой ключ.
    Если нет, то делает вставку с нужными данными, указанными в dict_kpi.
    Если есть, то делает вставку с нужными данными, указанными в dict_kpi и дублирует информацию из прошлых строк.

    @param dict_kpi: Словарь с описанием kpi.
    @return: Строку для вставки значений в SCD.
    """

    # Проверка наличия kpi_id в таблице
    df_check = pg.execute_to_df(f'''
    SELECT
        kpi_id
    FROM
        {TABLE}
    WHERE
        kpi_id = {dict_kpi['kpi_id']}
    ''')

    # Проверяем есть ли такой kpi_id в таблице
    if len(df_check) >= 1:
        # В запросе исключаем те поля, которые генерируется сами через `DEFAULT`
        query = f'''
            SELECT 
                column_name 
            FROM 
                information_schema.columns 
            WHERE 
                table_name = '{TABLE}'
                AND column_name NOT IN (
                    'id', 'created_at', 'changed_at', 'is_actual', 
                    {', '.join(f"'{i}'" for i in dict_kpi)}
                )
        '''

        df = pg.execute_to_df(query)  # noqa: PD901

        insert_sql_column_current = ', '.join(value for value in df.column_name)
        insert_sql_column_modified = insert_sql_column_current + f''', {', '.join(i for i in dict_kpi)}'''

        list_values = []

        for value in dict_kpi.values():
            # Обработка одинарных кавычек в значениях. Они встречаются при указании дат.
            if "'" in str(value):
                value = value.replace("'", "''")
                list_values.append(f"'{value}'")
            elif value is None:
                list_values.append('NULL')
            else:
                list_values.append(f"'{value}'")

        insert_sql_column_values = insert_sql_column_current + f''', {', '.join(list_values)}'''

        sql_insert = f'''
        INSERT INTO {TABLE}
        (
            {insert_sql_column_modified}
        )
        SELECT 
            {insert_sql_column_values} 
        FROM 
            {TABLE}
        WHERE
            is_actual IS TRUE
            AND kpi_id = {dict_kpi['kpi_id']};
        '''
    else:
        # Если нет такого kpi_id в таблице, то генерируем вставку значений из словаря
        columns = ', '.join(value for value in dict_kpi)

        list_values = []

        for value in dict_kpi.values():
            if "'" in str(value):
                value = value.replace("'", "''")
                list_values.append(f"'{value}'")
            elif value is None:
                list_values.append('NULL')
            else:
                list_values.append(f"'{value}'")

        values = ', '.join(list_values)

        sql_insert = f'''
            INSERT INTO {TABLE}({columns})
            VALUES ({values});
            '''

    return sql_insert


def scd_dim_kpi(dict_kpi: dict = None) -> None:
    """
    Основная функция, которая принимает на вход словарь с описанием kpi.

    Каждый ключ – это поле в таблице SCD.
    Каждое значение – это значение поля в таблице SCD.

    @param dict_kpi: Словарь с описанием kpi по выбранным колонкам.
    @return: Ничего не возвращает, выполняет SQL-скрипт на вставку данных в SCD.
    """

    # Обновление changed_at в предыдущей актуальной записи
    update_changed_at_for_kpi_id = f'''
    UPDATE {TABLE} 
    SET 
        changed_at = NOW() 
    WHERE 
        kpi_id = {dict_kpi['kpi_id']} 
        AND is_actual IS TRUE;
    '''

    # Вставка новой записи с обновленными значениями полей
    insert_new_values_for_kpi_id = gen_insert_sql_for_kpi_id(dict_kpi=dict_kpi)

    # Обновление is_actual для каждого kpi_id
    update_is_actual_for_kpi_id = f'''
    UPDATE {TABLE} 
    SET 
        is_actual = false 
    WHERE 
        kpi_id = {dict_kpi['kpi_id']} 
        AND id <> (
            SELECT MAX(id) 
            FROM {TABLE} 
            WHERE kpi_id = {dict_kpi['kpi_id']}
        );
    '''

    # Собираем SQL-скрипт из разных кусков, чтобы он прошел в одной транзакции
    sql_query = update_changed_at_for_kpi_id + insert_new_values_for_kpi_id + update_is_actual_for_kpi_id

    print(sql_query)  # noqa: T201
    pg.execute_script(sql_query)

from datetime import datetime, timedelta

import pendulum

from airflow import DAG

from airflow.sensors.external_task import ExternalTaskSensor

from airflow.operators.empty import EmptyOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator

from airflow.providers.postgres.hooks.postgres import PostgresHook


def create_kpi_of_dag(
        owner: str = None,
        dag_id: str = None,
        pg_target_schema: str = None,
        pg_target_table: str = None,
        index_kpi: str = None,
        pg_environment: str = None,
        airflow_environment: str = None,
        long_description: str = None,
        short_description: str = None,
        start_date: str = None,
        cron_expression: str = None,
        tags: str = None,
        sensors: str = None,
        sql_query: str = None,
) -> DAG:
    """
    Функция, которая генерирует типовой DAG для получения метрики.

    Вся логика описана и прокомментирована внутри функции.

    Некоторые моменты обработаны исключительно для функции, чтобы обработать какие-то атрибуты и получить желаемый
    эффект.

    @param owner: Владелец DAG.
    @param dag_id: Название DAG.
    @param pg_target_schema: Целевая схема в PostgreSQL.
    @param pg_target_table: Целевая таблица в PostgreSQL.
    @param index_kpi: ID показателя.
    @param pg_environment: Окружение PostgreSQL.
    @param airflow_environment: Окружение Airflow.
    @param long_description: Полное описание отчета.
    @param short_description: Короткое описание отчета.
    @param start_date: Дата начала работы DAG.
    @param cron_expression: Cron.
    @param tags: Tags.
    @param sensors: Sensors.
    @param sql_query: SQL-запрос для получения метрики.
    @return: Возвращает DAG.
    """
    # Конфигурация DAG
    local_tz = pendulum.timezone('Europe/Moscow')

    # Используемые таблицы в DAG
    pg_target_schema = pg_target_schema
    pg_target_table = pg_target_table
    pg_tmp_schema = 'stg'
    pg_tmp_table = f'tmp_{dag_id}_{{{{ data_interval_start.format("YYYY_MM_DD") }}}}'

    # Названия коннекторов к GP
    pg_connect = 'test_db_dev' if pg_environment == 'dev' else 'test_db'

    # Сделана заглушка атрибута. Это можно использовать для указания разных сценариев в зависимости от окружения
    airflow_environment = airflow_environment

    # Дата приходит в формате str и после парсинга, мы можем получить дату и любые элементы даты
    parse_date = pendulum.parse(start_date)

    args = {
        'owner': owner,
        'start_date': datetime(parse_date.year, parse_date.month, parse_date.day, tzinfo=local_tz),
        'catchup': True,
        'depends_on_past': True,
        'retries': 3,
        'retry_delay': timedelta(hours=1),
    }

    # Tags приходят в str формате, поэтому нужно их правильно "разобрать" и превратить в list
    raw_tags = list(tags.split(','))
    tags_ = []

    for i in raw_tags:
        tags_.append(  # noqa: PERF401
            i.replace("'", "")
            .replace(" ", '')
            .replace("[", "")
            .replace("]", "")
        )

    # Sensors приходят в str формате, поэтому нужно их правильно "разобрать" и превратить в list
    if sensors:
        raw_sensors = list(sensors.split(','))
        sensors_ = []
        for i in raw_sensors:
            sensors_.append(  # noqa: PERF401
                i.replace("'", "")
                .replace(' ', '')
                .replace("[", "")
                .replace("]", "")
            )
    else:
        sensors_ = None

    with DAG(
            dag_id=dag_id,
            schedule_interval=cron_expression,
            default_args=args,
            tags=tags_,
            description=short_description,
            concurrency=1,
            max_active_tasks=1,
            max_active_runs=1,
    ) as dag:
        dag.doc_md = long_description

        start = EmptyOperator(
            task_id='start',
        )

        # Если есть sensors, то мы создаем задачи с сенсорами, иначе создаем одну пустышку
        if sensors_:
            sensors_task = [
                ExternalTaskSensor(
                    task_id=f'sensor_{dag}',
                    external_dag_id=dag,
                    allowed_states=['success'],
                    mode='reschedule',
                    timeout=360000,  # длительность работы сенсора
                    poke_interval=600  # частота проверки
                ) for dag in sensors_
            ]
        else:
            sensors_task = [EmptyOperator(task_id=f'empty_{value}') for value in range(1)]

        drop_tmp_before = PostgresOperator(
            task_id='drop_tmp_before',
            sql=f'''DROP TABLE IF EXISTS {pg_tmp_schema}.{pg_tmp_table}''',
            postgres_conn_id=pg_connect
        )

        create_tmp = PostgresOperator(
            task_id='create_tmp',
            sql=f'''
            CREATE TABLE {pg_tmp_schema}.{pg_tmp_table} AS
            {
                sql_query.format(
                    start_date="{{ data_interval_start.format('YYYY-MM-DD') }}",
                    end_date="{{ data_interval_end.format('YYYY-MM-DD') }}"
                )
            };
            ''',
            postgres_conn_id=pg_connect
        )

        delete_from_target = PostgresOperator(
            task_id='delete_from_target',
            sql=f'''
            DELETE FROM {pg_target_schema}.{pg_target_table}
            WHERE 
                date IN (
                    SELECT 
                        date 
                    FROM 
                        {pg_tmp_schema}.{pg_tmp_table}
                    WHERE
                        kpi_id = {index_kpi}
                    )
            AND kpi_id = {index_kpi}
            ''',
            postgres_conn_id=pg_connect
        )

        insert_from_tmp_to_target = PostgresOperator(
            task_id='insert_from_tmp_to_target',
            sql=f'''
            INSERT INTO {pg_target_schema}.{pg_target_table}("date", value, kpi_id)
            SELECT 
                "date", 
                value, 
                {index_kpi} AS kpi_id 
            FROM 
                {pg_tmp_schema}.{pg_tmp_table}
            ''',
            postgres_conn_id=pg_connect
        )

        drop_tmp_after = PostgresOperator(
            task_id='drop_tmp_after',
            sql=f'''DROP TABLE IF EXISTS {pg_tmp_schema}.{pg_tmp_table}''',
            postgres_conn_id=pg_connect
        )

        end = EmptyOperator(
            task_id='end',
        )

        start >> sensors_task >> drop_tmp_before >> create_tmp >> delete_from_target >>\
        insert_from_tmp_to_target >> drop_tmp_after >> end

    return dag


# build a dag from dag config
def generator_of_morning_kpi_dag_to_gp() -> None:
    """
    Функция получает список config из БД и генерирует DAG's на основании функции `create_kpi_of_dag`.

    Итерируется по config и каждый раз выполняет функцию `create_kpi_of_dag`, которая возвращает DAG.

    @return: None
    """
    pg_hook = PostgresHook(postgres_conn_id='test_db')

    df = pg_hook.get_pandas_df(  # noqa: PD901
        '''
        SELECT 
            kpi_id,
            dag_id,
            "owner",
            sql_query,
            start_date,
            pg_environment,
            airflow_environment,
            short_description_md,
            long_description_md,
            cron,
            sensors,
            tags
        FROM 
            dim_kpi_dag_gen_config
        WHERE 
            is_actual IS TRUE 
        ORDER BY 
            id;
        '''
    )

    for i in range(len(df)):
        create_kpi_of_dag(
            owner=df.iloc[i].owner,
            dag_id=df.iloc[i].dag_id,
            pg_target_schema='public',
            pg_target_table='fct_dm_kpi',
            index_kpi=df.iloc[i].kpi_id,
            pg_environment=df.iloc[i].pg_environment,
            airflow_environment=df.iloc[i].airflow_environment,
            long_description=df.iloc[i].long_description_md,
            short_description=df.iloc[i].short_description_md,
            start_date=df.iloc[i].start_date,
            cron_expression=df.iloc[i].cron,
            tags=df.iloc[i].tags,
            sensors=df.iloc[i].sensors,
            sql_query=df.iloc[i].sql_query,
        )


generator_of_morning_kpi_dag_to_gp()

new_values = {
    'kpi_id': 1,
    'sql_query': '''
    SELECT 
        ('2023-'||((random()*11+1)::int)::varchar||'-'||((random()*27+1)::int)::varchar)::date AS date,
        (random()*100)::int AS value
    ''',
}


# Вызов функции
scd_dim_kpi(
    dict_kpi=new_values
)

# Пример использования
new_values = {
    'kpi_id': 6,
    'dag_id': 'test_6',
    'metric_name_en': 'test_6',
    'owner': 'korsak0v',
    'start_date': '2000-01-01',
    'cron': '15 14 1 * *',
    'tags': '''['dm', 'pg', 'gen_dag', 'from_pg']''',
    'sql_query': '''
    SELECT 
        ('2023-'||((random()*11+1)::int)::varchar||'-'||((random()*27+1)::int)::varchar)::date AS date,
        (random()*100)::int AS value
    ''',
}


# Вызов функции
scd_dim_kpi(
    dict_kpi=new_values
)

new_values = {
    'kpi_id': 7,
    'sensors': '''['test_1', 'test_2', 'test_3', 'test_4', 'test_5', 'test_6']''',
}


# Вызов функции
scd_dim_kpi(
    dict_kpi=new_values
)