Practice — Airflow DAG for End-to-End ELT — Data Pipeline Orchestration

This exercise builds a complete Airflow-style ELT pipeline for IPL match data using the pure-Python orchestration framework from Module 1 — enhanced with all Module 2 concepts: TaskFlow-style function chaining, XCom simulation, dynamic task mapping via `expand`, sensor behaviour, branching, and idempotent Delta-style writes. The pipeline extracts match metadata and delivery records, waits for the official scorecard sensor, transforms and validates the data, dynamically processes each match independently, and loads results using idempotent upserts with retry.

The exercise is structured in three steps. Step 1 builds the extraction and sensor layer, verifying that the sensor correctly blocks downstream tasks until its condition is met and returns to normal flow when the condition becomes true. Step 2 implements dynamic match processing with simulated `expand` and verifies per-match retry isolation. Step 3 implements the full pipeline end-to-end with a branching quality gate, idempotency test, and final assertions on row counts, retry counts, and DLQ routing.

Analogy🏏Cricket

🏏 Think of it like cricket: Migrating from Airflow to Prefect is like the same bowling coach shifting from traditional Test cricket notation to a modern T20 analytics dashboard — the underlying ball-by-ball data (the business logic) is exactly the same. What changes is how the data is recorded, displayed, and acted upon. The yorker that Bumrah bowls in over 20 is identical whether it is recorded in the old scorebook (Airflow DAG file) or the new analytics platform (Prefect flow). The migration is a transcription exercise, not a strategy change — and a wise coach verifies that the runs, wickets, and economies match exactly between the old and new system before decommissioning the scorebook. That verification step is the whole heart of the migration: because the yorker is unchanged, the only honest test is to run the same over through both systems and confirm the recorded runs, wickets and economies match to the last digit before the old scorebook is thrown away. Rushing to burn the scorebook the moment the shiny dashboard lights up is how teams lose a season of records to a silent transcription slip. The coach keeps both systems running in parallel for a while, reconciles their outputs ball by ball, and only when every figure agrees does he trust the new dashboard alone — a transcription is only complete when you have proven nothing was lost in the copying.

Step 1 — Extraction, Sensor, and XCom Layer

Build the first three pipeline stages: a mock S3 sensor that blocks until a scorecard file 'arrives', an extract task that fetches match metadata and delivery records, and an XCom layer that passes the discovered match IDs to the downstream dynamic processing stage. Verify the sensor correctly blocks downstream tasks when the file is not yet available and unblocks them when the condition is met. Assert that the XCom store contains the correct match ID list after the extract task completes.

Analogy🏏Cricket

🏏 Think of it like cricket: This stage is the start of the official post-match data routine. The mock S3 sensor is the scorer refusing to begin tabulation until the signed scorecard physically lands on the desk — poll after poll he checks, and only when the file has genuinely arrived does he let the downstream work begin, exactly as the sensor blocks tasks until its condition is met and unblocks them once it turns true. The extract task is the scorer reading off the match metadata and delivery records, and the XCom layer is the handwritten note he passes to the next desk listing precisely which fixtures are in play today — MI vs CSK, RCB vs KKR, KKR vs SRH. Just as the next statistician must receive that exact list of matches before he can share the work out, the downstream dynamic stage depends on the XCom store holding the correct match ID list. Verifying that the sensor gates correctly and the note carries the right IDs is what keeps the whole day's processing from ever starting on incomplete or wrong data.

python

# exercise_airflow_elt.py — Step 1: Extraction, sensor, and XCom layer
from datetime import date, timedelta
from collections import defaultdict
from typing import Optional, Any
import time
import random
import pandas as pd
import numpy as np

np.random.seed(42)
random.seed(42)

# ── Mock infrastructure ───────────────────────────────────────────────────────
xcom_store:   dict = {}
warehouse_db: dict[str, list[dict]] = defaultdict(list)
dlq_store:    list = []
retry_counts: dict = defaultdict(int)

class MockSensor:
    """
    Simulates an S3KeySensor.
    Polls a condition function at poke_interval; raises TimeoutError at timeout.
    """
    def __init__(self, task_id: str, condition_fn,
                 poke_interval_s: float = 0.05, timeout_s: float = 1.0,
                 mode: str = "reschedule"):
        self.task_id        = task_id
        self.condition_fn   = condition_fn
        self.poke_interval  = poke_interval_s
        self.timeout        = timeout_s
        self.mode           = mode   # reschedule releases slot between polls

    def run(self) -> bool:
        start = time.monotonic()
        attempt = 0
        while True:
            attempt += 1
            if self.condition_fn():
                print(f"  ✓ {self.task_id} — condition met after {attempt} poll(s)")
                return True
            if (time.monotonic() - start) > self.timeout:
                raise TimeoutError(f"{self.task_id} timed out after {self.timeout}s")
            time.sleep(self.poke_interval)

# ── S3 mock: scorecard file "arrives" after 2 polls ──────────────────────────
s3_files: set = set()
poll_count = [0]

def scorecard_available() -> bool:
    poll_count[0] += 1
    if poll_count[0] >= 2:   # file "arrives" on 2nd poll
        s3_files.add("scorecards/2024-04-20/scorecard.csv")
    return "scorecards/2024-04-20/scorecard.csv" in s3_files

scorecard_sensor = MockSensor(
    "wait_for_scorecard", scorecard_available,
    poke_interval_s=0.05, timeout_s=2.0, mode="reschedule"
)

# Test 1: sensor eventually succeeds
sensor_result = scorecard_sensor.run()
assert sensor_result == True
assert poll_count[0] >= 2
print(f"  Sensor polled {poll_count[0]} time(s) before condition met ✓")

# ── Extract task ───────────────────────────────────────────────────────────────
def extract_matches(logical_date: date, **_) -> list[dict]:
    matches = [
        {"match_id": 10001, "venue": "Wankhede",     "home": "MI",  "away": "CSK"},
        {"match_id": 10002, "venue": "Chinnaswamy",  "home": "RCB", "away": "KKR"},
        {"match_id": 10003, "venue": "Eden Gardens", "home": "KKR", "away": "SRH"},
    ]
    xcom_store[("extract_matches", "return_value")] = matches
    xcom_store[("extract_matches", "match_ids")]    = [m["match_id"] for m in matches]
    return matches

def extract_deliveries(logical_date: date, **_) -> dict[int, list[dict]]:
    deliveries_by_match = {}
    for match_id in [10001, 10002, 10003]:
        np.random.seed(match_id)
        deliveries_by_match[match_id] = [
            {"delivery_id": i, "match_id": match_id,
             "runs": int(np.random.choice([0,1,2,4,6])),
             "is_wicket": bool(np.random.random() < 0.05)}
            for i in range(1, 121)  # 120 deliveries per match
        ]
    xcom_store[("extract_deliveries", "return_value")] = deliveries_by_match
    return deliveries_by_match

LOGICAL_DATE = date(2024, 4, 20)

matches    = extract_matches(LOGICAL_DATE)
deliveries = extract_deliveries(LOGICAL_DATE)

# Verify XCom
assert xcom_store[("extract_matches", "match_ids")] == [10001, 10002, 10003]
assert len(xcom_store[("extract_deliveries", "return_value")]) == 3
print(f"  XCom: {len(matches)} matches, "
      f"{sum(len(v) for v in deliveries.values())} deliveries ✓")

Step 2 — Dynamic Match Processing and Quality Gate

Implement the dynamic match processing stage using `simulate_expand` from Module 1, where each match is processed by an independent task instance with its own retry budget. Inject a transient failure on the first attempt for one specific match and verify that it retries and succeeds while the other matches are unaffected. Follow with the quality gate branch: if all matches have null-venue rate below 5%, route to the warehouse loader; otherwise route to the DLQ. Run the full pipeline and assert all invariants.

Analogy🏏Cricket

🏏 Think of it like cricket: Dynamic match processing is the head scorer assigning one independent statistician to each fixture rather than making a single official grind through them in sequence. Each statistician carries his own retry budget — if the data feed for RCB vs KKR drops out for a moment, that one scorer simply re-does his match while the MI and KKR statisticians' finished work stands untouched, exactly as the injected transient failure on match 10002 retries in isolation without disturbing the others. Then the quality gate is the head statistician's sign-off before publication: if the null-venue rate stays below the threshold the figures go to the official warehouse, but if too many venues are missing the batch is diverted to the review queue instead. Just as one match's late feed must never hold up every other match's statistics, isolating retries per fixture and branching bad data to the DLQ is what keeps the clean matches flowing while the problem ones are quarantined for manual review.

python

# exercise_airflow_elt.py — Step 2: Dynamic processing, branching, idempotency
from datetime import date

# ── Simulate expand(): per-match task instances ───────────────────────────────
fail_once_match = {10002}   # match 10002 fails on first attempt

def process_match(
    match_id: int, logical_date: date,
    matches: list[dict], deliveries: dict,
    retries: int = 2,
) -> dict:
    """Process one match — called once per mapped instance."""
    retry_counts[match_id] += 1
    if match_id in fail_once_match and retry_counts[match_id] == 1:
        fail_once_match.discard(match_id)  # only fail once
        raise ConnectionError(f"Transient error fetching match {match_id}")

    match_meta  = next(m for m in matches if m["match_id"] == match_id)
    match_deliv = deliveries[match_id]
    merged = [{**d, **match_meta} for d in match_deliv]
    return {"match_id": match_id, "rows": len(merged),
            "venue": match_meta["venue"], "status": "success"}

def expand_with_retry(match_ids, logical_date, matches, deliveries, retries=2):
    """Simulate Airflow expand() with per-instance retry."""
    results = []
    for mid in match_ids:
        for attempt in range(retries + 1):
            try:
                result = process_match(mid, logical_date, matches, deliveries)
                print(f"  ✓ process_match[{mid}] — success (attempt {attempt+1})")
                results.append(result)
                break
            except Exception as exc:
                if attempt < retries:
                    print(f"  ↻ process_match[{mid}] — retry {attempt+1}: {exc}")
                else:
                    print(f"  ✗ process_match[{mid}] — failed after {retries+1} attempts")
                    results.append({"match_id": mid, "status": "failed", "rows": 0})
    return results

results = expand_with_retry(
    [10001, 10002, 10003], LOGICAL_DATE, matches, deliveries
)
assert all(r["status"] == "success" for r in results)
assert retry_counts[10002] == 2   # failed once, succeeded on retry
assert retry_counts[10001] == 1   # never failed
print(f"  Dynamic processing: {len(results)} matches, "
      f"match 10002 required {retry_counts[10002]} attempt(s) ✓")

# ── Quality gate branch ───────────────────────────────────────────────────────
def quality_gate(results: list[dict]) -> str:
    """Branch operator: returns next task_id based on quality check."""
    failed = [r for r in results if r["status"] != "success"]
    if failed:
        return "route_to_dlq"
    # Check null venue rate
    null_venue_pct = sum(1 for r in results if not r.get("venue")) / len(results)
    return "load_to_warehouse" if null_venue_pct < 0.05 else "route_to_dlq"

branch_target = quality_gate(results)
print(f"  Quality gate → {branch_target}")
assert branch_target == "load_to_warehouse"

# ── Load to warehouse (idempotent) ───────────────────────────────────────────
def load_to_warehouse(results: list[dict], logical_date: date) -> int:
    warehouse_db[str(logical_date)] = results
    return len(results)

def route_to_dlq(results: list[dict], logical_date: date) -> int:
    dlq_store.extend([{**r, "logical_date": str(logical_date)} for r in results])
    return len(results)

if branch_target == "load_to_warehouse":
    loaded = load_to_warehouse(results, LOGICAL_DATE)
else:
    route_to_dlq(results, LOGICAL_DATE)
    loaded = 0

# Idempotency: run the full load again for same logical_date
loaded_again = load_to_warehouse(results, LOGICAL_DATE)
wh_count_run1 = len(warehouse_db[str(LOGICAL_DATE)])
wh_count_run2 = len(warehouse_db[str(LOGICAL_DATE)])  # same dict entry overwritten
assert wh_count_run1 == wh_count_run2 == 3
print(f"  Idempotency: warehouse has {wh_count_run1} match records (stable) ✓")

# Final summary
print("\n=== Pipeline Summary ===")
print(f"  Logical date:    {LOGICAL_DATE}")
print(f"  Matches loaded:  {len(warehouse_db[str(LOGICAL_DATE)])}")
print(f"  DLQ events:      {len(dlq_store)}")
print(f"  Retry counts:    {dict(retry_counts)}")
print("  All assertions passed. Airflow ELT pipeline exercise complete.")

Warning: In production Airflow, the `BranchPythonOperator` marks all non-selected downstream tasks as `skipped`, not `failed`. Any task with the default `trigger_rule=all_success` that is downstream of skipped tasks will also be marked `upstream_failed` and never run. To ensure a summary or cleanup task always runs after a branch, set `trigger_rule=TriggerRule.ALL_DONE` — this runs the task regardless of whether upstream tasks are `success`, `failed`, or `skipped`. Missing this on cleanup tasks is the most common branching-related bug in production Airflow DAGs.

Extension Challenge: Extend the pipeline with a final summary task (using `trigger_rule=ALL_DONE`) that runs after both the `load_to_warehouse` and `route_to_dlq` branches, reads the warehouse and DLQ counts, and publishes a Slack-style summary message. The summary should include: total matches processed, total rows loaded, DLQ count, any matches that required retries, and the logical date. This tests the `ALL_DONE` trigger rule and the post-branch aggregation pattern that production pipelines use for end-of-run reporting.