طراحی پایپلاینهای دادهای بازگشتپذیر (Idempotent) و رهگیری متادیتا با Airflow و OpenLineage

from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.exceptions import AirflowFailException

DEFAULT_ARGS = {
    'owner': 'data-platform',
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
}

@dag(
    dag_id='fin_daily_settlement_wap',
    schedule='@daily',
    start_date=datetime(2025, 1, 1),
    catchup=True,
    max_active_runs=3,
    default_args=DEFAULT_ARGS,
    tags=['finance', 'idempotent', 'wap']
)
def financial_settlement_pipeline():

    @task
    def stage_partition(data_interval_start=None, data_interval_end=None) -> str:
        dt_str = data_interval_start.strftime('%Y_%m_%d')
        staging_table = f'staging.settlements_{dt_str}'
        hook = PostgresHook(postgres_conn_id='warehouse_postgres')

        sql = f"""
        DROP TABLE IF EXISTS {staging_table};
        CREATE TABLE {staging_table} (LIKE prod.daily_settlements INCLUDING ALL);
        INSERT INTO {staging_table}
        SELECT
            DATE('{data_interval_start.strftime('%Y-%m-%d')}') AS settlement_date,
            merchant_id,
            COUNT(tx_id) AS tx_count,
            SUM(amount_cents) AS total_amount_cents
        FROM raw.payment_transactions
        WHERE created_at >= '{data_interval_start.isoformat()}'
          AND created_at < '{data_interval_end.isoformat()}'
          AND status = 'CAPTURED'
        GROUP BY 1, 2;
        """
        hook.run(sql)
        return staging_table

    @task
    def audit_staging_data(staging_table: str) -> str:
        hook = PostgresHook(postgres_conn_id='warehouse_postgres')
        row_count, negative_totals = hook.get_first(f"""
            SELECT
                COUNT(*),
                COALESCE(SUM(CASE WHEN total_amount_cents < 0 THEN 1 ELSE 0 END), 0)
            FROM {staging_table};
        """)
        if row_count == 0:
            raise AirflowFailException(f'Data Quality Audit Failed: {staging_table} is empty!')
        if negative_totals > 0:
            raise AirflowFailException(f'Audit Failed: Found {negative_totals} negative settlement rows!')
        return staging_table

    @task
    def publish_atomic_swap(staging_table: str, data_interval_start=None, data_interval_end=None) -> None:
        dt_str = data_interval_start.strftime('%Y_%m_%d')
        start_date = data_interval_start.strftime('%Y-%m-%d')
        end_date = data_interval_end.strftime('%Y-%m-%d')
        target_partition = f'prod.daily_settlements_p_{dt_str}'
        hook = PostgresHook(postgres_conn_id='warehouse_postgres')

        # Atomic metadata swap inside a single transaction block
        swap_sql = f"""
        BEGIN;
        CREATE TABLE IF NOT EXISTS {target_partition}
            PARTITION OF prod.daily_settlements
            FOR VALUES FROM ('{start_date}') TO ('{end_date}');
        ALTER TABLE prod.daily_settlements DETACH PARTITION {target_partition};
        DROP TABLE {target_partition};
        ALTER TABLE {staging_table} RENAME TO daily_settlements_p_{dt_str};
        ALTER TABLE prod.daily_settlements.daily_settlements_p_{dt_str} SET SCHEMA prod;
        ALTER TABLE prod.daily_settlements ATTACH PARTITION prod.daily_settlements_p_{dt_str}
            FOR VALUES FROM ('{start_date}') TO ('{end_date}');
        COMMIT;
        """
        hook.run(swap_sql)

    staged = stage_partition()
    audited = audit_staging_data(staged)
    publish_atomic_swap(audited)

dag_instance = financial_settlement_pipeline()