Practice — Design a DAG for a Daily ETL Workflow

This exercise applies every concept from Module 1 by designing and implementing a complete DAG for the IPL daily ETL workflow — from raw API extraction through transformation, quality checks, and warehouse loading. You will build the DAG using the pure-Python `DAG` class from Lesson 1, implement idempotent task functions scoped to a logical date, configure retry with exponential backoff, add event-driven S3 file detection via a mock sensor, and verify the topological sort produces the correct execution order before running any task.

The pipeline processes one day of IPL match data: extract match metadata from a REST API and delivery records from a database, wait for the official scorecard file to land in S3, transform and merge both sources, run data quality checks, load to the analytics warehouse, and trigger a downstream dbt model run. The exercise ends with an idempotency verification: the full pipeline runs twice for the same logical date and the warehouse row count assertion confirms no duplication occurred.

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 — DAG Structure and Task Functions

Build the full DAG with eight tasks: `wait_for_scorecard` (sensor), `extract_api_matches`, `extract_db_deliveries`, `validate_sources`, `transform_and_merge`, `check_data_quality`, `load_to_warehouse`, and `trigger_dbt`. Define all task functions with the `logical_date` parameter as their idempotency key, and assign sensible retry counts to extract tasks. Verify the topological sort produces a valid execution order respecting all declared dependencies by running a dry run before executing any actual task logic.

Analogy🏏Cricket

🏏 Think of it like cricket: Building this DAG's structure is like a captain naming the exact batting order and roles before a single ball is bowled. The sensor is the twelfth man watching for the covers to come off — no play until the pitch is genuinely ready. The extract tasks are your openers, and just as you give an opener extra time to settle after a shaky first ball, you assign retry counts to those extract tasks so a transient failure earns another delivery rather than a wicket. Just as every batsman is told to build on the same match situation, every task takes logical_date as its idempotency key so the whole innings is anchored to one fixed date. And just as a captain walks the order on paper before the toss — confirming no one bats out of turn — the dry-run topological sort confirms the execution order respects every declared dependency before any real task runs. Naming the order, roles and retries up front is what lets the innings unfold cleanly instead of collapsing mid-over.

python

# exercise_dag_design.py — Step 1: DAG structure and task functions
from datetime import date, datetime, timezone
from collections import defaultdict, deque
from typing import Callable, Any, Optional
import time
import random
import pandas as pd
import numpy as np

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

# ── Reuse the DAG + Task classes from Lesson 1 ───────────────────────────────
class Task:
    def __init__(self, task_id: str, fn: Callable,
                 retries: int = 0, retry_delay_s: float = 1.0,
                 trigger_rule: str = "all_success"):
        self.task_id       = task_id
        self.fn            = fn
        self.retries       = retries
        self.retry_delay_s = retry_delay_s
        self.trigger_rule  = trigger_rule
        self.upstream: list['Task'] = []

    def __rshift__(self, other: 'Task') -> 'Task':
        other.upstream.append(self)
        return other

class DAGRunner:
    def __init__(self, tasks: list[Task]):
        self.tasks   = {t.task_id: t for t in tasks}
        self.results: dict[str, dict] = {}

    def _topo_order(self) -> list[Task]:
        in_degree = {tid: 0 for tid in self.tasks}
        for t in self.tasks.values():
            for u in t.upstream:
                in_degree[t.task_id] += 1
        q     = deque(tid for tid, d in in_degree.items() if d == 0)
        order = []
        while q:
            tid = q.popleft()
            order.append(self.tasks[tid])
            for t in self.tasks.values():
                if self.tasks[tid] in t.upstream:
                    in_degree[t.task_id] -= 1
                    if in_degree[t.task_id] == 0:
                        q.append(t.task_id)
        assert len(order) == len(self.tasks), "Cycle detected"
        return order

    def run(self, logical_date: date, dry_run: bool = False) -> dict:
        print(f"\n{'DRY RUN — ' if dry_run else ''}Running DAG for {logical_date}")
        print("-" * 55)
        for task in self._topo_order():
            # Check trigger rule
            upstream_states = [
                self.results.get(u.task_id, {}).get("state", "pending")
                for u in task.upstream
            ]
            skip = (
                task.trigger_rule == "all_success"
                and any(s != "success" for s in upstream_states)
            )
            if skip:
                self.results[task.task_id] = {"state": "upstream_failed", "output": None}
                print(f"  ⊘ {task.task_id:<30} [upstream_failed]")
                continue

            if dry_run:
                self.results[task.task_id] = {"state": "success", "output": f"<dry_{task.task_id}>"}
                print(f"  ~ {task.task_id:<30} [dry_run]")
                continue

            # Run with retry
            upstream_ctx = {
                u.task_id: self.results[u.task_id]["output"]
                for u in task.upstream
            }
            for attempt in range(1, task.retries + 2):
                try:
                    output = task.fn(logical_date=logical_date, **upstream_ctx)
                    self.results[task.task_id] = {"state": "success", "output": output}
                    print(f"  ✓ {task.task_id:<30} [success]")
                    break
                except Exception as exc:
                    if attempt <= task.retries:
                        delay = task.retry_delay_s * (2 ** (attempt - 1))
                        print(f"  ↻ {task.task_id:<30} [retry {attempt}/{task.retries}] — {exc}")
                        time.sleep(min(delay, 0.1))  # cap at 100ms for exercise speed
                    else:
                        self.results[task.task_id] = {"state": "failed", "output": None}
                        print(f"  ✗ {task.task_id:<30} [failed] — {exc}")
        return self.results

# ── Idempotent task function implementations ──────────────────────────────────

WAREHOUSE: dict[str, list[dict]] = defaultdict(list)  # in-memory warehouse

def wait_for_scorecard(logical_date: date, **_) -> bool:
    """Mock sensor: immediately finds the scorecard file."""
    print(f"    Checking for scorecard: s3://ipl-data/scorecards/{logical_date}.csv")
    return True  # file found

def extract_api_matches(logical_date: date, **_) -> list[dict]:
    matches = [
        {"match_id": 10001, "venue": "Wankhede",    "date": str(logical_date)},
        {"match_id": 10002, "venue": "Chinnaswamy",  "date": str(logical_date)},
    ]
    return matches

def extract_db_deliveries(logical_date: date, **_) -> list[dict]:
    np.random.seed(int(logical_date.strftime("%Y%m%d")))
    return [
        {"delivery_id": i, "match_id": 10001 + (i // 120),
         "runs": int(np.random.choice([0,1,2,4,6]))}
        for i in range(1, 241)   # 2 matches × 120 deliveries
    ]

def validate_sources(logical_date: date,
                     wait_for_scorecard, extract_api_matches, extract_db_deliveries,
                     **_) -> dict:
    assert wait_for_scorecard == True
    assert len(extract_api_matches)  > 0
    assert len(extract_db_deliveries) > 0
    return {"matches": len(extract_api_matches), "deliveries": len(extract_db_deliveries)}

def transform_and_merge(logical_date: date,
                        extract_api_matches, extract_db_deliveries, **_) -> pd.DataFrame:
    matches_df   = pd.DataFrame(extract_api_matches)
    deliveries_df= pd.DataFrame(extract_db_deliveries)
    merged = deliveries_df.merge(matches_df, on="match_id", how="left")
    merged["logical_date"] = str(logical_date)
    return merged

def check_data_quality(logical_date: date, transform_and_merge: pd.DataFrame, **_) -> dict:
    assert transform_and_merge["runs"].between(0, 6).all(), "Invalid runs values"
    assert transform_and_merge["venue"].notna().all(), "Null venues after join"
    return {"rows": len(transform_and_merge), "quality": "pass"}

def load_to_warehouse(logical_date: date,
                      transform_and_merge: pd.DataFrame, check_data_quality, **_) -> int:
    """Idempotent: delete existing rows for this logical_date, then insert."""
    WAREHOUSE[str(logical_date)] = transform_and_merge.to_dict(orient="records")
    return len(WAREHOUSE[str(logical_date)])

def trigger_dbt(logical_date: date, load_to_warehouse, **_) -> str:
    print(f"    Triggering dbt run for {logical_date} ({load_to_warehouse} rows loaded)")
    return "dbt_run_complete"

print("Task functions defined")

Step 2 — Wire Dependencies and Run Idempotency Test

Wire all eight tasks with correct upstream dependencies, confirm the topological sort with a dry run, execute the full pipeline twice for the same logical date, and assert that the warehouse row count is identical after both runs — confirming idempotency. Simulate a transient failure on the first attempt of `extract_api_matches` and verify the retry logic recovers correctly before the DAG proceeds to downstream tasks.

Analogy🏏Cricket

🏏 Think of it like cricket: Wiring the dependencies with `>>` is like confirming no batsman takes strike before the one ahead of him is out or off the field — the sensor must clear before extraction, extraction before the merge, the merge before the load. Just as a captain runs a quick shadow-practice of the order before the toss, the dry run confirms the topological sort is valid without facing a real ball. The idempotency test is the board's trust check: replay the exact same match data twice and confirm the official scorecard shows an identical row count both times, because a load that duplicated rows on rerun could never be trusted for a live final. And just as a coach deliberately signals a dropped catch in rehearsal to confirm the backup fielder covers, injecting a transient failure on the first extract attempt confirms the retry logic recovers before the innings continues. Practising the wiring and the reruns is what proves the pipeline scores the same clean total every single evening.

python

# exercise_dag_design.py — Step 2: Wire dependencies and idempotency test
from datetime import date

LOGICAL_DATE = date(2024, 4, 20)

# Build tasks
t_sensor   = Task("wait_for_scorecard",   wait_for_scorecard,   retries=2)
t_api      = Task("extract_api_matches",  extract_api_matches,  retries=3, retry_delay_s=0.5)
t_db       = Task("extract_db_deliveries",extract_db_deliveries,retries=2)
t_validate = Task("validate_sources",     validate_sources)
t_transform= Task("transform_and_merge",  transform_and_merge)
t_quality  = Task("check_data_quality",   check_data_quality)
t_load     = Task("load_to_warehouse",    load_to_warehouse)
t_dbt      = Task("trigger_dbt",          trigger_dbt)

# Wire dependencies
t_sensor   >> t_api
t_sensor   >> t_db
t_api      >> t_validate
t_db       >> t_validate
t_validate >> t_transform
t_transform>> t_quality
t_quality  >> t_load
t_load     >> t_dbt

all_tasks = [t_sensor, t_api, t_db, t_validate, t_transform, t_quality, t_load, t_dbt]
runner    = DAGRunner(all_tasks)

# Dry run: verify topological order without executing
print("Topological order (dry run):")
dry_results = runner.run(LOGICAL_DATE, dry_run=True)

# Simulate transient failure on first attempt of extract_api_matches
call_count = [0]
original_extract = extract_api_matches
def flaky_extract(logical_date, **kw):
    call_count[0] += 1
    if call_count[0] == 1:
        raise ConnectionError("Simulated transient API timeout")
    return original_extract(logical_date=logical_date, **kw)

t_api.fn = flaky_extract
runner2  = DAGRunner(all_tasks)

# Run 1: first attempt fails, retry succeeds
results1 = runner2.run(LOGICAL_DATE)
assert results1["load_to_warehouse"]["state"] == "success"
wh_count_run1 = len(WAREHOUSE[str(LOGICAL_DATE)])

# Run 2: idempotency test — same logical date, re-run entire DAG
t_api.fn = extract_api_matches   # restore non-flaky version
runner3  = DAGRunner(all_tasks)
results3 = runner3.run(LOGICAL_DATE)
assert results3["load_to_warehouse"]["state"] == "success"
wh_count_run2 = len(WAREHOUSE[str(LOGICAL_DATE)])

print(f"\nIdempotency test:")
print(f"  Run 1 warehouse rows: {wh_count_run1}")
print(f"  Run 2 warehouse rows: {wh_count_run2}")
assert wh_count_run1 == wh_count_run2, \
    f"Idempotency failed: {wh_count_run1} → {wh_count_run2} rows after second run"
print(f"  Row count unchanged ✓ — pipeline is idempotent")

# Final summary
print("\nFinal task states (Run 2):")
for task_id, result in results3.items():
    icon = {"success":"✓", "failed":"✗", "upstream_failed":"⊘"}.get(result["state"],"?")
    print(f"  {icon} {task_id:<30} [{result['state']}]")
print("\nAll assertions passed. DAG design exercise complete.")

Warning: The task functions in this exercise receive upstream task outputs as keyword arguments by task ID — `transform_and_merge` receives `extract_api_matches` and `extract_db_deliveries` as parameters. In production Airflow, task outputs are passed via XCom (small values) or shared storage paths (large DataFrames). Never pass a full DataFrame through XCom — it is serialised to the metadata database and can easily exceed the 48KB size limit, causing the task to fail at runtime. Pass file paths or storage references between tasks; read the data from storage inside each task.

Extension Challenge: Extend the pipeline with a branching task after `check_data_quality` that takes two paths: if the quality check reports more than 5% null venues (indicating a bad API response), it routes to a `quarantine_to_s3` task that writes the bad data to an S3 quarantine prefix and sends an alert; otherwise it routes to `load_to_warehouse`. Implement this as a conditional task that reads the quality check result and sets downstream task states accordingly. This is the orchestration equivalent of the Airflow `BranchPythonOperator`.