Practice — Add GE Checks to an Airflow Pipeline

This exercise integrates Great Expectations quality checks into the IPL daily ELT pipeline from Module 2, adding a two-layer quality gate: a pre-transformation GE checkpoint that validates the raw extracted data, and a post-transformation check that validates the merged output before loading. The exercise also adds freshness checking, a tiered alerting function that routes critical failures to a mock PagerDuty endpoint and warnings to a mock Slack webhook, and an OpenLineage event emitter. The pipeline only proceeds to the load stage if both quality gates pass.

The exercise is structured in three steps. Step 1 builds the GE suite and integrates the pre-load quality gate, verifying that the pipeline halts when bad data is injected. Step 2 adds freshness checking, SLA monitoring, and tiered alerting, verifying that different failure types route to the correct alert channel. Step 3 adds the OpenLineage instrumentation and runs the complete end-to-end pipeline with assertions on quality gate results, alert routing, and lineage event structure.

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 — GE Suite Integration and Quality Gate

Build the two GE suites — one for raw delivery records (completeness, validity, uniqueness checks) and one for the merged output (venue not null, economy in range, row count conservation). Run the pipeline with clean data to verify both gates pass, then inject a known bad row (delivery with runs=9) and verify the pre-transformation gate halts the pipeline before the transform stage executes, with the failure details available in the quality result summary.

Analogy🏏Cricket

🏏 Think of it like cricket: These two GE suites are the two-stage verification the official scorers run before any statistic is published. The raw-delivery suite is the ball-by-ball check on the scorecard as it arrives — every delivery has an ID, a match, a runs value, and crucially no delivery records an impossible 9 off the bat when only 0 through 6 are legal. The merged suite is the audit after the innings is assembled: venue present on every row, economy inside a sane range, no balls lost or gained in the join. Just as a scorer who spots a runs=9 entry halts and refuses to fold that scorecard into the season totals, injecting that bad row makes the pre-transformation gate stop the pipeline before the transform stage ever runs. Practising both the clean pass and the deliberate bad-row halt is what proves the gate genuinely blocks corrupt data at the source rather than letting it silently reach the official record.

python

# exercise_ge_pipeline.py — Step 1: GE suite integration and quality gates
import great_expectations as gx
import pandas as pd
import numpy as np
from datetime import date, datetime, timezone, timedelta
from typing import Optional

np.random.seed(42)

# ── GE quality gate factory ─────────────────────────────────────────────────
def build_raw_delivery_suite(context, suite_name="raw_delivery_suite"):
    suite = context.add_expectation_suite(suite_name)
    datasource = context.sources.add_pandas("ipl_raw")
    asset = datasource.add_dataframe_asset("raw_deliveries")
    v = context.get_validator(
        batch_request=asset.build_batch_request(),
        expectation_suite=suite,
    )
    v.expect_column_values_to_not_be_null("delivery_id")
    v.expect_column_values_to_not_be_null("match_id")
    v.expect_column_values_to_not_be_null("runs")
    v.expect_column_values_to_be_between("runs", min_value=0, max_value=6)
    v.expect_column_values_to_be_unique("delivery_id")
    v.expect_table_row_count_to_be_between(min_value=1, max_value=10_000)
    v.save_expectation_suite(discard_failed_expectations=False)
    return suite_name

def run_quality_gate(
    df: pd.DataFrame,
    suite_name: str,
    context,
    stage: str,
    halt_on_fail: bool = True,
) -> dict:
    """Run a GE suite against df. Raise if halt_on_fail and suite fails."""
    asset = context.sources.get_datasource(f"ipl_{stage}").assets[0] \
        if f"ipl_{stage}" in [s.name for s in context.list_datasources()] \
        else context.sources.add_pandas(f"ipl_{stage}").add_dataframe_asset(f"{stage}_df")
    batch_req  = asset.build_batch_request(dataframe=df)
    validator  = context.get_validator(
        batch_request=batch_req,
        expectation_suite_name=suite_name,
    )
    result     = validator.validate()
    failures   = [{"check": r.expectation_config.expectation_type,
                   "col":   r.expectation_config.kwargs.get("column","table")}
                  for r in result.results if not r.success]
    summary = {
        "stage":    stage,
        "success":  bool(result.success),
        "checks":   len(result.results),
        "failures": failures,
    }
    print(f"  Quality gate [{stage}]: {summary['success']} "
          f"({sum(1 for r in result.results if r.success)}/{summary['checks']} passed)")
    if not summary["success"] and halt_on_fail:
        raise ValueError(f"[{stage}] Quality gate FAILED: {failures}")
    return summary

# ── Pipeline with quality gates ──────────────────────────────────────────────
WAREHOUSE = {}

def run_pipeline_with_quality(
    logical_date: date,
    inject_bad_row: bool = False,
) -> dict:
    context   = gx.get_context(mode="ephemeral")
    suite_raw = build_raw_delivery_suite(context)

    # Extract
    deliveries = pd.DataFrame({
        "delivery_id":  range(1, 241),
        "match_id":     [10001]*120 + [10002]*120,
        "runs":         np.random.choice([0,1,2,4,6], 240),
        "bowler":       np.random.choice(["Bumrah","Shami"], 240),
        "batting_team": np.random.choice(["Mumbai Indians","CSK"], 240),
        "updated_at":   [datetime.now(timezone.utc) - timedelta(hours=1)] * 240,
    })
    matches = pd.DataFrame({
        "match_id": [10001, 10002],
        "venue":    ["Wankhede", "Chinnaswamy"],
    })

    if inject_bad_row:
        bad = pd.DataFrame([{"delivery_id": 9999, "match_id": 10001,
                              "runs": 9, "bowler": "Unknown",
                              "batting_team": "CSK",
                              "updated_at": datetime.now(timezone.utc)}])
        deliveries = pd.concat([deliveries, bad], ignore_index=True)

    # Pre-transform quality gate
    gate_raw = run_quality_gate(deliveries, suite_raw, context, "raw")
    if not gate_raw["success"]:
        return {"halted_at": "pre_transform", "reason": gate_raw["failures"]}

    # Transform
    merged = deliveries.merge(matches, on="match_id", how="left")

    # Load
    WAREHOUSE[str(logical_date)] = merged.to_dict(orient="records")
    return {"success": True, "rows": len(merged)}

# Run 1: clean data
result1 = run_pipeline_with_quality(date(2024, 4, 20))
print(f"  Clean run: {result1}")
assert result1.get("success") == True

# Run 2: inject bad row — pipeline should halt at pre_transform gate
result2 = run_pipeline_with_quality(date(2024, 4, 21), inject_bad_row=True)
print(f"  Bad data run: {result2}")
assert result2.get("halted_at") == "pre_transform"
assert date(2024, 4, 21).isoformat() not in WAREHOUSE  # nothing written
print("Step 1 ✓: quality gate halted pipeline on bad data")

Step 2 — Freshness, SLA and Tiered Alerts

Add a freshness check that fails if the delivery data is older than 4 hours, an SLA monitor that tracks pipeline start time and alerts if the run exceeds 10 minutes, and a tiered alert router that sends critical failures to a mock PagerDuty endpoint and warnings to a mock Slack webhook. Verify each alert type fires correctly by testing three scenarios: a fresh-data pipeline run (no alerts), a stale-data run (freshness warning to Slack), and a quality gate failure (critical alert to PagerDuty).

Analogy🏏Cricket

🏏 Think of it like cricket: This step adds the match referee's clocks and graded flags to the statistics pipeline. The freshness check is the submission deadline — if the delivery data is more than four hours stale it is treated like a scorecard filed too late to trust. The SLA monitor is the over-rate clock: a run that drags past ten minutes draws a flag even though it eventually finishes, just as a slow over-rate is penalised even when no wicket falls. The tiered router is the graded response — a genuinely dangerous integrity breach goes straight to the senior official on the critical PagerDuty channel, while a minor staleness warning is a quiet note on the Slack channel — matching the severity of the intervention to the severity of the problem. Practising all three scenarios — fresh, stale and failed — confirms each signal reaches the right channel, so the serious alarms are never drowned out by the trivial ones.

python

# exercise_ge_pipeline.py — Step 2: Freshness, SLA, tiered alerting
from datetime import date, datetime, timezone, timedelta
import pandas as pd
import numpy as np

np.random.seed(42)

# ── Mock alert sinks ─────────────────────────────────────────────────────────
pagerduty_alerts: list = []
slack_messages:   list = []

def mock_pagerduty(title: str, details: str) -> None:
    pagerduty_alerts.append({"title": title, "details": details,
                              "ts": datetime.now(timezone.utc).isoformat()})
    print(f"  [PAGERDUTY] {title}")

def mock_slack(message: str, severity: str = "warning") -> None:
    slack_messages.append({"message": message, "severity": severity,
                            "ts": datetime.now(timezone.utc).isoformat()})
    print(f"  [SLACK/{severity.upper()}] {message}")

def tiered_alert(severity: str, title: str, details: str) -> None:
    if severity == "critical":
        mock_pagerduty(title, details)
    elif severity == "warning":
        mock_slack(details, severity="warning")
    else:
        mock_slack(details, severity="info")

# ── Freshness check ──────────────────────────────────────────────────────────
def check_freshness_gate(
    df: pd.DataFrame,
    ts_col: str,
    max_age_hours: float,
) -> dict:
    latest = pd.to_datetime(df[ts_col]).max()
    if latest.tzinfo is None:
        latest = latest.tz_localize("UTC")
    age = (datetime.now(timezone.utc) - latest).total_seconds() / 3600
    passed = age <= max_age_hours
    return {"passed": passed, "age_hours": round(age, 2), "max_age_hours": max_age_hours}

# ── SLA Monitor ──────────────────────────────────────────────────────────────
import time

class SLAMonitor:
    def __init__(self, max_minutes: float):
        self.max_minutes = max_minutes
        self.start_time  = datetime.now(timezone.utc)

    def check(self) -> dict:
        elapsed = (datetime.now(timezone.utc) - self.start_time).total_seconds() / 60
        return {"elapsed_min": round(elapsed, 2),
                "breach":      elapsed > self.max_minutes,
                "max_min":     self.max_minutes}

# ── Test three scenarios ─────────────────────────────────────────────────────
pagerduty_alerts.clear()
slack_messages.clear()

# Scenario 1: fresh data — no alerts
deliveries_fresh = pd.DataFrame({
    "delivery_id": range(1, 121),
    "runs":       np.random.choice([0,1,2,4,6], 120),
    "updated_at": [datetime.now(timezone.utc) - timedelta(hours=1)] * 120,
})
freshness1 = check_freshness_gate(deliveries_fresh, "updated_at", max_age_hours=4)
assert freshness1["passed"] == True
print(f"  Scenario 1 (fresh): {freshness1['age_hours']}h < {freshness1['max_age_hours']}h ✓")

# Scenario 2: stale data — Slack warning
deliveries_stale = deliveries_fresh.copy()
deliveries_stale["updated_at"] = datetime.now(timezone.utc) - timedelta(hours=6)
freshness2 = check_freshness_gate(deliveries_stale, "updated_at", max_age_hours=4)
assert freshness2["passed"] == False
tiered_alert("warning",
    "IPL data freshness warning",
    f"Deliveries data is {freshness2['age_hours']}h old (max {freshness2['max_age_hours']}h)")
assert len(slack_messages) == 1
print(f"  Scenario 2 (stale): Slack warning fired ✓")

# Scenario 3: quality gate failure — PagerDuty critical
tiered_alert("critical",
    "IPL pipeline CRITICAL: quality gate failed",
    "Pre-transform GE check failed: runs=9 violates accepted range [0,6]")
assert len(pagerduty_alerts) == 1
print(f"  Scenario 3 (quality fail): PagerDuty alert fired ✓")

# ── SLA test: simulate a fast run within SLA
sla = SLAMonitor(max_minutes=10)
time.sleep(0.01)  # simulate pipeline work
sla_check = sla.check()
print(f"  SLA check: {sla_check['elapsed_min']}min (limit {sla_check['max_min']}min) — "
      f"{'BREACH' if sla_check['breach'] else 'OK'}")
assert not sla_check["breach"]
print("Step 2 ✓: freshness, SLA, and tiered alerting verified")

Warning: GE's ephemeral context recreates all datasources on every call, which means a call to `build_raw_delivery_suite` inside `run_pipeline_with_quality` also creates a new datasource object. If you call the function twice in the same context, the second call will raise a `DatasourceAlreadyExistsError`. Either reuse the context object across calls (passing it as a parameter) or wrap the datasource creation in a `try/except` that catches the duplicate error and loads the existing datasource instead. The exercise above creates a new context per pipeline run to avoid this — in production, reuse a module-level context.

Extension Challenge: Add a third step that simulates the OpenLineage instrumentation from Lesson 22. Wrap the `run_pipeline_with_quality` function with START and COMPLETE/FAIL lineage events using the `emit_lineage_event` function, passing the raw deliveries source and the warehouse destination as input and output datasets. Run the pipeline successfully and print the emitted lineage event JSON to verify it contains the correct job name, namespace, input datasets, and output datasets. This connects all four Module 4 concepts: quality dimensions, GE checks, freshness monitoring, and lineage emission.