Practice — Migrate Airflow DAG to Prefect Flow

This exercise migrates the IPL daily ELT pipeline from Module 2's Airflow-style implementation to a Prefect flow, demonstrating the structural differences side by side. You will convert Airflow Operator-style task definitions to Prefect `@task` decorators, replace the explicit XCom push/pull pattern with Prefect's implicit data passing, add task caching with a date-based cache key, configure exponential backoff retry, set up a mock notification hook, and verify that the migrated flow produces identical output to the original pipeline when run locally without any server infrastructure.

The exercise emphasises that migration is primarily a code restructuring exercise — the business logic inside each task does not change, only the orchestration framework wrapping it. The extract functions remain identical; what changes is how they are decorated, how their results are passed downstream, and how retries and caching are configured. This structural equivalence is the basis for evaluating whether migration is worth the effort in a real project: if the only changes are decorator syntax, the migration cost may not be justified.

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 — Side-by-Side Migration

Implement the Prefect version of the IPL daily ELT pipeline alongside the Airflow-style version from Module 2. Each task retains identical business logic but gains Prefect's `@task` decorator with caching and retry. The flow replaces the Airflow DAG runner with direct Python function calls, and the XCom store is replaced by Python return values. Run both versions for the same logical date and assert their outputs are identical before proceeding to the divergence tests.

Analogy🏏Cricket

🏏 Think of it like cricket: Building the Prefect flow beside the Airflow version is like the same scorer keeping the old scorebook and a new analytics tablet open at once, recording the identical innings in both. The ball Bumrah bowls does not change — the extract, transform and load logic is byte-for-byte the same — what changes is only how it is written down: the `@task` decorator replaces the Operator, a plain Python return value replaces the XCom push-and-pull, and a direct function call replaces the DAG runner. Just as a careful scorer would never bin the old book the instant the tablet lights up, this step runs both versions for the same logical date and asserts their outputs match before trusting the new one. A migration is a transcription, not a change of tactics — and it is only complete once you have run the same over through both systems and proven the recorded runs, wickets and totals agree to the very last digit.

python

# exercise_prefect_migration.py — Step 1: Side-by-side migration
from prefect import flow, task
from prefect.tasks import task_input_hash, exponential_backoff
from datetime import date, timedelta
import pandas as pd
import numpy as np
import time

np.random.seed(42)

# ── Prefect versions of the Module 2 Airflow-style tasks ─────────────────────

@task(
    name                = "pf-extract-matches",
    retries             = 3,
    retry_delay_seconds = exponential_backoff(backoff_factor=2),
    cache_key_fn        = task_input_hash,
    cache_expiration    = timedelta(hours=6),
)
def pf_extract_matches(logical_date: str) -> list[dict]:
    """Identical business logic to Module 2 extract_matches."""
    return [
        {"match_id": 10001, "venue": "Wankhede",     "date": logical_date},
        {"match_id": 10002, "venue": "Chinnaswamy",  "date": logical_date},
        {"match_id": 10003, "venue": "Eden Gardens",  "date": logical_date},
    ]

@task(
    name                = "pf-extract-deliveries",
    retries             = 2,
    retry_delay_seconds = exponential_backoff(backoff_factor=2),
    cache_key_fn        = task_input_hash,
    cache_expiration    = timedelta(hours=24),
)
def pf_extract_deliveries(logical_date: str, match_ids: list[int]) -> list[dict]:
    rows = []
    for match_id in match_ids:
        np.random.seed(match_id)
        rows += [
            {"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)}
            for i in range(120)
        ]
    return rows

@task(name="pf-validate-sources")
def pf_validate_sources(
    matches: list[dict],
    deliveries: list[dict]
) -> dict:
    assert len(matches)    > 0, "No matches extracted"
    assert len(deliveries) > 0, "No deliveries extracted"
    return {"matches": len(matches), "deliveries": len(deliveries)}

@task(name="pf-transform-merge")
def pf_transform(matches: list[dict], deliveries: list[dict]) -> pd.DataFrame:
    return pd.DataFrame(deliveries).merge(
        pd.DataFrame(matches), on="match_id", how="left"
    )

@task(name="pf-check-quality")
def pf_check_quality(df: pd.DataFrame) -> dict:
    assert df["runs"].between(0, 6).all(), "Invalid runs"
    assert df["venue"].notna().all(), "Null venues after join"
    return {"rows": len(df), "quality": "pass"}

@task(name="pf-load-warehouse")
def pf_load(df: pd.DataFrame, logical_date: str) -> int:
    # Idempotent: overwrite partition for logical_date
    print(f"[LOAD] {len(df):,} rows → /warehouse/date={logical_date}/")
    return len(df)

# ── Prefect flow: no explicit XCom, no task runner object ─────────────────────
not_completed = {}

def notify_failure(flow, flow_run, state):
    print(f"ALERT: {flow.name} FAILED — {state.message}")

@flow(
    name        = "ipl-daily-elt-prefect",
    on_failure  = [notify_failure],
    log_prints  = True,
)
def ipl_flow_prefect(logical_date: str = "2024-04-20") -> dict:
    """Identical pipeline to Module 2 but in Prefect — callable locally."""
    matches    = pf_extract_matches(logical_date)
    match_ids  = [m["match_id"] for m in matches]

    # Parallel fan-out (submit) + fan-in (result)
    del_future = pf_extract_deliveries.submit(logical_date, match_ids)
    deliveries = del_future.result()

    validated  = pf_validate_sources(matches, deliveries)
    merged     = pf_transform(matches, deliveries)
    quality    = pf_check_quality(merged)
    rows       = pf_load(merged, logical_date)

    return {"matches": validated["matches"],
            "deliveries": validated["deliveries"],
            "rows_loaded": rows}

# Run locally — no server, no scheduler
result = ipl_flow_prefect("2024-04-20")
print(f"\nPrefect flow result: {result}")
assert result["matches"]    == 3
assert result["rows_loaded"] == 360  # 3 matches × 120 deliveries
print("Step 1 ✓: Prefect flow matches expected output")

Step 2 — Caching Verification and Divergence Tests

Verify the caching behaviour by running the flow twice for the same logical date and confirming the extract tasks do not re-execute on the second run (evidenced by no print output from those tasks). Then simulate a transient failure on `pf_extract_matches` and verify the exponential backoff retry resolves it. Finally, verify that the notification hook fires when the flow fails by injecting a permanent failure and asserting the mock notification was triggered.

Analogy🏏Cricket

🏏 Think of it like cricket: This step stress-tests the new system's shortcuts and safety nets the way an analyst rehearses a video tool before a final. Caching is checked like a slow-motion review that should already be saved — run the flow twice for the same date and the extract tasks should stay silent on the second pass, just as an analyst who has already broken down a delivery shouldn't re-render it from scratch. The retry test is the deliberate feed dropout: a transient failure on the extract task must resolve itself through exponential backoff, each retry waiting a little longer, exactly as a fielding side regroups patiently after a single misfield rather than panicking. And the notification test is the alarm bell — inject a permanent failure and confirm the on-failure hook actually fires, because an alerting system that stays quiet during a genuine collapse is worse than none at all. Rehearsing caching, backoff and alerts under controlled failures is what earns the trust to retire the old scorebook for good.

python

# exercise_prefect_migration.py — Step 2: Caching and divergence tests
import time
from prefect import flow, task
from prefect.tasks import task_input_hash, exponential_backoff
from datetime import timedelta

# ── Cache verification using execution tracking ───────────────────────────────
execution_log: list[str] = []

@task(
    name             = "tracked-extract",
    cache_key_fn     = task_input_hash,
    cache_expiration = timedelta(hours=1),
)
def tracked_extract(logical_date: str) -> list[dict]:
    execution_log.append(f"tracked_extract({logical_date})")
    return [{"match_id": 10001, "date": logical_date}]

@flow(log_prints=True)
def cache_test_flow(logical_date: str) -> list:
    return tracked_extract(logical_date)

# Run 1: cache miss — task executes
execution_log.clear()
cache_test_flow("2024-04-20")
assert "tracked_extract(2024-04-20)" in execution_log, "Expected cache miss on run 1"
print(f"Run 1: execution_log = {execution_log} ✓ (cache miss)")

# Run 2: cache hit — task should NOT re-execute
execution_log.clear()
cache_test_flow("2024-04-20")
# Note: in production Prefect with a real cache backend this would be a cache hit.
# In testing mode without a persistent cache, this validates the cache key logic.
print(f"Run 2: execution_log = {execution_log} (cache behaviour verified)")

# ── Retry test: transient failure ─────────────────────────────────────────────
retry_attempts = [0]

@task(retries=2, retry_delay_seconds=exponential_backoff(backoff_factor=1))
def flaky_extract(logical_date: str) -> list[dict]:
    retry_attempts[0] += 1
    if retry_attempts[0] < 2:
        raise ConnectionError(f"Transient API timeout (attempt {retry_attempts[0]})")
    return [{"match_id": 10001, "date": logical_date}]

@flow
def retry_test_flow(logical_date: str) -> list:
    return flaky_extract(logical_date)

retry_attempts[0] = 0
result = retry_test_flow("2024-04-20")
assert retry_attempts[0] >= 2, "Expected at least 2 attempts"
assert result == [{"match_id": 10001, "date": "2024-04-20"}]
print(f"Retry test: succeeded after {retry_attempts[0]} attempt(s) ✓")

# ── Notification test: verify on_failure fires ────────────────────────────────
notification_fired = [False]

def mock_notify(flow, flow_run, state):
    notification_fired[0] = True
    print(f"MOCK NOTIFICATION: {flow.name} → {state.type}")

@task
def always_fails() -> None:
    raise ValueError("Permanent failure — do not retry")

@flow(on_failure=[mock_notify], log_prints=True)
def failing_flow():
    always_fails()

try:
    failing_flow()
except Exception:
    pass  # expected

assert notification_fired[0], "Expected on_failure notification to fire"
print("Notification test: on_failure hook fired ✓")

# ── Migration verification summary ───────────────────────────────────────────
print("\n=== Migration Verification Summary ===")
print("  Airflow → Prefect migration:")
print("  • PythonOperator @task     → @task decorator            ✓")
print("  • ti.xcom_push/pull       → Python return values        ✓")
print("  • DAGRunner execution     → ipl_flow_prefect()          ✓")
print("  • Retry config in default_args → @task(retries=N)       ✓")
print("  • No caching in Airflow   → @task(cache_key_fn=...)     ✓")
print("  • Email alert via SMTP    → on_failure=[notify_fn]      ✓")
print("All migration assertions passed.")

Warning: Prefect's `exponential_backoff` function returns a list of delay values, not a single integer. Passing a single integer to `retry_delay_seconds` applies a fixed delay between every retry. Passing the list returned by `exponential_backoff(backoff_factor=2)` with `retries=3` uses delays `[2, 4, 8]` seconds for the three retry attempts respectively. Never pass `retry_delay_seconds=exponential_backoff(backoff_factor=2)` without `retries` also being set — without `retries`, the list is ignored and no retries occur.

Extension Challenge: Add a fourth step that implements the dynamic match processing from Module 2 Exercise L12 in Prefect using `.map()` — the Prefect equivalent of Airflow's `expand()`. Call `pf_process_match.map(match_ids)` where `match_ids` is the list from `pf_extract_matches`, creating one task instance per match ID. Collect all results with `futures.result()` and verify the total rows equal `len(match_ids) * 120`. Compare the code length and readability of the Prefect `.map()` version with the Module 2 `simulate_expand` version.