This exercise builds the orchestration layer of the capstone pipeline using the pure-Python DAG framework from Module 1: a mock S3 sensor, dynamic task mapping that creates one extract task per match ID, a GE-powered branching quality gate that routes clean data to the warehouse loader and bad data to the DLQ, and a cleanup task with `ALL_DONE` trigger rule. The pipeline runs for a logical date, verifies all task states, and asserts idempotency by running twice and confirming the warehouse row count is identical after both runs.
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.
The mock infrastructure simulates the full Airflow execution model — sensor polling, XCom-style result passing between tasks, dynamic task instance creation, branching with skip propagation, and retry on transient failures — without requiring a running Airflow cluster. This approach ensures the pipeline logic can be tested and verified in CI before deploying to a real Airflow environment, which is the production pattern for any organisation that practices test-driven pipeline development.
Step 1 — Sensor and Dynamic Extraction
Implement the S3 sensor that blocks until the scorecard file is available, the `get_match_ids` discovery task that returns the list of match IDs for the logical date, and the dynamic `extract_match` task that runs once per match ID with retry on transient failures. Verify that the sensor polls at least twice before succeeding, that exactly three match tasks are created for the three match IDs in the test data, and that a transient failure on one match is retried while the other matches process successfully in parallel.
Analogy🏏Cricket
🏏 Think of it like cricket: The S3 sensor that blocks until the scorecard file arrives is the third umpire refusing to let the next over begin until the ball is confirmed to have crossed the rope. He does nothing dramatic — he simply keeps looking, poll after poll, at whether the condition is finally true, and only when the file has genuinely landed does he wave play on to the extraction tasks. Then dynamic task mapping takes over like a head scorer who, on seeing how many matches are being played today, assigns one independent scorer per fixture rather than making a single overworked official cover them all in sequence. Each mapped extract task is its own scorer with its own scorecard, so if the feed for match 10002 fails, only that one scorer re-does his match while the others' completed work stands untouched. The sensor gates the start; the mapping parallelises the work — together they turn a rigid, wait-then-loop-through-everything routine into a disciplined start-on-signal, then-scatter-the-scorers operation.
python
# capstone_airflow_dag.py — Step 1: Sensor and dynamic extraction
from datetime import date, datetime, timezone, timedelta
from collections import defaultdict
from typing import Any, Optional, Callable
import time
import random
import numpy as np
import pandas as pd
np.random.seed(42)
random.seed(42)
# ── Mock infrastructure (from Module 2 exercise, extended) ────────────────────
xcom_store: dict = {}
warehouse_db: dict = {}
dlq_store: list = []
lineage_log: list = []
poll_count = [0]
class MockSensor:
def __init__(self, task_id, condition_fn,
poke_interval=0.05, timeout=1.0, mode="reschedule"):
self.task_id = task_id
self.condition_fn = condition_fn
self.poke_interval= poke_interval
self.timeout = timeout
self.mode = mode
def run(self) -> bool:
start = time.monotonic()
while True:
poll_count[0] += 1
if self.condition_fn():
print(f" ✓ {self.task_id} — condition met after {poll_count[0]} poll(s)")
return True
if (time.monotonic() - start) > self.timeout:
raise TimeoutError(f"{self.task_id} timed out")
time.sleep(self.poke_interval)
# S3 scorecard "arrives" on 2nd poll
s3_files = set()
def scorecard_available():
if poll_count[0] >= 2:
s3_files.add("scorecards/2024-04-20/scorecard.csv")
return "scorecards/2024-04-20/scorecard.csv" in s3_files
sensor = MockSensor("wait_for_scorecard", scorecard_available)
sensor_result = sensor.run()
assert sensor_result == True and poll_count[0] >= 2
print(f" Sensor polls: {poll_count[0]} ✓")
# ── get_match_ids: discovery task ─────────────────────────────────────────────
def get_match_ids(logical_date: date, **_) -> list[int]:
xcom_store[("get_match_ids", "return_value")] = [10001, 10002, 10003]
return [10001, 10002, 10003]
match_ids = get_match_ids(date(2024, 4, 20))
assert match_ids == [10001, 10002, 10003]
print(f" get_match_ids: {match_ids} ✓")
# ── Dynamic extract: one task per match, with retry ──────────────────────────
MATCHES_META = {
10001: {"match_id": 10001, "venue": "Wankhede", "home": "MI", "away": "CSK"},
10002: {"match_id": 10002, "venue": "Chinnaswamy", "home": "RCB", "away": "KKR"},
10003: {"match_id": 10003, "venue": "Eden Gardens", "home": "KKR", "away": "SRH"},
}
retry_counts = defaultdict(int)
fail_once = {10002} # match 10002 fails on first attempt
def extract_match(match_id: int, logical_date: date, retries: int = 2) -> dict:
retry_counts[match_id] += 1
if match_id in fail_once and retry_counts[match_id] == 1:
fail_once.discard(match_id)
raise ConnectionError(f"Transient API timeout for match {match_id}")
np.random.seed(match_id)
deliveries = [
{"delivery_id": match_id * 1000 + i,
"match_id": match_id,
"runs": int(np.random.choice([0,1,2,4,6])),
"is_wicket": bool(np.random.random() < 0.05),
"bowler": np.random.choice(["Bumrah","Shami","Hardik"]),
"updated_at": (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat()}
for i in range(120)
]
return {"match_id": match_id, "meta": MATCHES_META[match_id],
"deliveries": deliveries}
def expand_with_retry(fn, items, retries=2, **kwargs):
results = []
for item in items:
for attempt in range(retries + 1):
try:
result = fn(item, **kwargs)
results.append(result)
if attempt > 0:
print(f" ↻ extract_match[{item}] — succeeded on attempt {attempt+1}")
break
except Exception as exc:
if attempt < retries:
pass
else:
results.append({"match_id": item, "status": "failed",
"error": str(exc)})
return results
extract_results = expand_with_retry(
extract_match, match_ids, retries=2, logical_date=date(2024, 4, 20)
)
assert len(extract_results) == 3
assert all("deliveries" in r or r.get("status") == "failed"
for r in extract_results)
assert retry_counts[10002] == 2 # failed once, succeeded on retry
assert retry_counts[10001] == 1 # never failed
print(f" Dynamic extraction: {len(extract_results)} matches, "
f"match 10002 required {retry_counts[10002]} attempt(s) ✓")
# ── GE quality gate (branching) ───────────────────────────────────────────────
import great_expectations as gx
all_deliveries = pd.DataFrame([
d for r in extract_results if "deliveries" in r
for d in r["deliveries"]
])
context = gx.get_context(mode="ephemeral")
ds = context.sources.add_pandas("cap_src")
asset = ds.add_dataframe_asset("deliveries")
batch_req = asset.build_batch_request(dataframe=all_deliveries)
validator = context.get_validator(
batch_request=batch_req,
expectation_suite=context.add_expectation_suite("cap_suite")
)
validator.expect_column_values_to_not_be_null("delivery_id")
validator.expect_column_values_to_be_between("runs", min_value=0, max_value=6)
validator.expect_column_values_to_be_unique("delivery_id")
validator.expect_table_row_count_to_be_between(min_value=1, max_value=10_000)
validator.save_expectation_suite(discard_failed_expectations=False)
ge_result = validator.validate()
branch_target = "load_to_warehouse" if ge_result.success else "route_to_dlq"
print(f" GE quality gate: success={ge_result.success} → {branch_target} ✓")
assert branch_target == "load_to_warehouse"
xcom_store[("quality_gate", "result")] = {
"success": bool(ge_result.success),
"rows": len(all_deliveries),
}
assert xcom_store[("quality_gate", "result")]["success"] == True
print("Step 1 ✓: sensor, dynamic extraction, GE quality gate complete")Step 2 — Load, Cleanup and Idempotency
Implement the warehouse loader with idempotent delete-and-reload semantics, the DLQ router for failed quality gate runs, and the `ALL_DONE` cleanup task. Run the complete pipeline twice for the same logical date and assert the warehouse row count is identical after both runs. Simulate a quality gate failure by injecting a bad row and verify the pipeline routes to DLQ rather than the warehouse, with the warehouse row count unchanged.
Analogy🏏Cricket
🏏 Think of it like cricket: This step is the scorer finalising the day's record and proving it can be safely redone. The idempotent loader uses delete-and-reload, exactly like a scorer correcting a day's entry by overwriting that day's line rather than adding a duplicate beneath it — so reprocessing April 20th replaces April 20th's rows instead of doubling them, and the row count stays identical across two runs. The DLQ router is the review pavilion: when a bad row fails the quality gate the batch is diverted there and the official warehouse is left untouched. The ALL_DONE cleanup task is the ground staff rolling the covers whether the innings ended in a century or a collapse — it must run even when the branch took the failure path. Practising the rerun, the DLQ diversion, and the always-fires cleanup is what guarantees a mid-day reprocess after fixing an issue never corrupts or duplicates the settled record.
python
# capstone_airflow_dag.py — Step 2: Load, cleanup, idempotency
from datetime import date
import pandas as pd
import numpy as np
np.random.seed(42)
LOGICAL_DATE = date(2024, 4, 20)
# ── Load and DLQ tasks ────────────────────────────────────────────────────────
def load_to_warehouse(deliveries_df: pd.DataFrame, logical_date: date) -> int:
"""Idempotent: delete existing rows for logical_date, then insert."""
warehouse_db[str(logical_date)] = deliveries_df.to_dict(orient="records")
return len(warehouse_db[str(logical_date)])
def route_to_dlq(deliveries_df: pd.DataFrame, reason: str,
logical_date: date) -> None:
dlq_store.append({
"logical_date": str(logical_date),
"reason": reason,
"rows": len(deliveries_df),
})
print(f" DLQ: {len(deliveries_df)} rows routed for '{reason}'")
def cleanup(logical_date: date, **_) -> str:
"""ALL_DONE: always runs — cleans staging files and temp resources."""
print(f" Cleanup complete for {logical_date}")
return "cleaned"
# ── Run 1: clean data ─────────────────────────────────────────────────────────
loaded1 = load_to_warehouse(all_deliveries, LOGICAL_DATE)
cleanup(LOGICAL_DATE)
wh_count_run1 = len(warehouse_db[str(LOGICAL_DATE)])
assert wh_count_run1 == 360 # 3 matches × 120 deliveries
print(f" Run 1: {wh_count_run1} rows loaded ✓")
# ── Run 2: idempotency — same logical date, should produce same row count ─────
# Re-extract and re-load
retry_counts_r2 = defaultdict(int)
fail_once_r2 = set() # no failures this time
def extract_match_clean(match_id: int, logical_date: date, **_) -> dict:
np.random.seed(match_id)
deliveries = [
{"delivery_id": match_id * 1000 + i,
"match_id": match_id,
"runs": int(np.random.choice([0,1,2,4,6])),
"is_wicket": bool(np.random.random() < 0.05),
"bowler": np.random.choice(["Bumrah","Shami","Hardik"]),
"updated_at": (datetime(2024,4,20,1,0,0).isoformat())}
for i in range(120)
]
return {"match_id": match_id, "deliveries": deliveries}
results2 = [extract_match_clean(mid, LOGICAL_DATE) for mid in [10001, 10002, 10003]]
deliveries2 = pd.DataFrame([
d for r in results2 for d in r["deliveries"]
])
loaded2 = load_to_warehouse(deliveries2, LOGICAL_DATE)
wh_count_run2 = len(warehouse_db[str(LOGICAL_DATE)])
assert wh_count_run1 == wh_count_run2, \
f"Idempotency failed: {wh_count_run1} → {wh_count_run2}"
print(f" Run 2: {wh_count_run2} rows (stable — idempotency ✓)")
# ── Quality gate failure → DLQ routing ───────────────────────────────────────
bad_row = pd.DataFrame([{"delivery_id": 9999, "match_id": 10001,
"runs": 9, "is_wicket": False,
"bowler": "Unknown",
"updated_at": "2024-04-20T01:00:00"}])
bad_df = pd.concat([all_deliveries, bad_row], ignore_index=True)
context2 = gx.get_context(mode="ephemeral")
ds2 = context2.sources.add_pandas("bad_src")
asset2 = ds2.add_dataframe_asset("bad_del")
batch_req2 = asset2.build_batch_request(dataframe=bad_df)
validator2 = context2.get_validator(
batch_request=batch_req2,
expectation_suite=context2.add_expectation_suite("bad_suite")
)
validator2.expect_column_values_to_be_between("runs", min_value=0, max_value=6)
validator2.save_expectation_suite(discard_failed_expectations=False)
bad_result = validator2.validate()
assert not bad_result.success
branch_bad = "route_to_dlq" if not bad_result.success else "load_to_warehouse"
assert branch_bad == "route_to_dlq"
route_to_dlq(bad_df, "runs_out_of_range", LOGICAL_DATE)
wh_after_bad = len(warehouse_db.get(str(LOGICAL_DATE), []))
assert wh_after_bad == wh_count_run2 # warehouse unchanged
cleanup(LOGICAL_DATE)
print(f" DLQ route: warehouse rows unchanged at {wh_after_bad} ✓")
print(f" DLQ events: {len(dlq_store)} ✓")
print("Step 2 ✓: load, cleanup, idempotency, DLQ routing complete")Warning: In real Airflow, a branching task that marks one path as skipped causes all tasks on the skipped path — including cleanup tasks — to be marked `upstream_failed` unless they have `trigger_rule=ALL_DONE`. This is the most common bug in Airflow DAGs that use branching: the cleanup task appears to run correctly in development where the quality gate always passes, but silently never runs in production when the DLQ path is taken. Always set `trigger_rule=TriggerRule.ALL_DONE` on cleanup tasks and verify the behaviour by explicitly triggering the failure path in your tests.
Extension Challenge: Extend the DAG with a `send_run_summary` task (also with `trigger_rule=ALL_DONE`) that reads the XCom result from the quality gate and warehouse loader to compose a summary message: 'IPL daily pipeline for 2024-04-20: 3 matches, 360 rows loaded, 0 DLQ events, cleanup complete'. Include the logical date, match count, row count, DLQ count, and cleanup status. This end-of-run summary is the standard pattern for automated pipeline health reporting that stakeholders receive each morning.