S050 · Orchestration — Airflow, DAGs, Retries, Backfills

# dags/orders_pipeline.py — end-to-end idempotent daily pipeline.
# Requires: Airflow >= 2.7. Uses TaskFlow, sensors, retries, SLAs.
from __future__ import annotations
 
import datetime as dt
import logging
from airflow.decorators import dag, task
from airflow.sensors.external_task import ExternalTaskSensor
from airflow.exceptions import AirflowSkipException
from airflow.providers.postgres.hooks.postgres import PostgresHook
 
log = logging.getLogger(__name__)
 
DEFAULT_ARGS = {
    "owner": "data-eng",
    "retries": 3,
    "retry_delay": dt.timedelta(minutes=5),
    "retry_exponential_backoff": True,
    "max_retry_delay": dt.timedelta(minutes=30),
    "sla": dt.timedelta(hours=2),
    "email_on_failure": True,
}
 
@dag(
    dag_id="orders_pipeline",
    default_args=DEFAULT_ARGS,
    schedule="0 3 * * *",              # 3am daily
    start_date=dt.datetime(2024, 1, 1),
    catchup=False,                       # don't backfill on new DAG
    max_active_runs=1,                   # never run two copies of same DAG at once
    tags=["orders", "daily", "warehouse"],
)
def orders_pipeline():
 
    # --------------------------------------------------------------
    # 0. Wait for upstream ingest DAG to finish for the same logical date
    # --------------------------------------------------------------
    wait_ingest = ExternalTaskSensor(
        task_id="wait_ingest_orders",
        external_dag_id="ingest_orders",
        external_task_id="publish_to_lake",
        allowed_states=["success"],
        failed_states=["failed", "upstream_failed", "skipped"],
        poke_interval=60,
        timeout=60 * 60 * 4,             # 4 h max wait
        mode="reschedule",               # release worker slot while waiting
    )
 
    # --------------------------------------------------------------
    # 1. Idempotent stage → INSERT OVERWRITE partition by execution_date
    # --------------------------------------------------------------
    @task(task_id="stage_orders")
    def stage_orders(logical_date: dt.datetime | None = None) -> int:
        ds = logical_date.date()
        hook = PostgresHook(postgres_conn_id="warehouse")
        # Deleting first + inserting = idempotent on re-run
        hook.run(f"""
            BEGIN;
            DELETE FROM stage.orders WHERE order_date = '{ds}';
            INSERT INTO stage.orders (order_id, customer_id, amount, order_date)
            SELECT order_id, customer_id, amount, order_date::date
              FROM lake.raw_orders
             WHERE order_date::date = '{ds}';
            COMMIT;
        """)
        (count,) = hook.get_first(
            "SELECT COUNT(*) FROM stage.orders WHERE order_date = %s", parameters=[ds]
        )
        log.info("staged %s rows for %s", count, ds)
        return count
 
    # --------------------------------------------------------------
    # 2. Transform → MERGE into mart.fact_orders (upsert on natural key)
    # --------------------------------------------------------------
    @task(task_id="load_fact")
    def load_fact(row_count: int, logical_date: dt.datetime | None = None) -> None:
        if row_count == 0:
            raise AirflowSkipException("no rows for this date")
        ds = logical_date.date()
        hook = PostgresHook(postgres_conn_id="warehouse")
        hook.run(f"""
            MERGE INTO mart.fact_orders t
            USING stage.orders s ON t.order_id = s.order_id
            WHEN MATCHED THEN UPDATE SET
              customer_id = s.customer_id, amount = s.amount, order_date = s.order_date
            WHEN NOT MATCHED THEN INSERT (order_id, customer_id, amount, order_date)
              VALUES (s.order_id, s.customer_id, s.amount, s.order_date);
        """)
 
    # --------------------------------------------------------------
    # 3. Data-quality gate (simple assertion; real teams use dbt tests / GE)
    # --------------------------------------------------------------
    @task(task_id="dq_check")
    def dq_check(logical_date: dt.datetime | None = None) -> None:
        ds = logical_date.date()
        hook = PostgresHook(postgres_conn_id="warehouse")
        (bad,) = hook.get_first(
            "SELECT COUNT(*) FROM mart.fact_orders WHERE order_date = %s AND amount < 0",
            parameters=[ds],
        )
        if bad > 0:
            raise ValueError(f"{bad} negative-amount rows found for {ds}")
 
    n = stage_orders()
    fact = load_fact(n)
    qc = dq_check()
    wait_ingest >> n
    fact >> qc
 
orders_pipeline()