Title here
Summary here
Lesson 4 - The Load Task
Welcome to the Load Task
Ingest, validate, and transform all worked entirely with files on disk. Load is the one task in this pipeline that talks to an outside system — which is exactly why Module 4 spent a whole module on hooks and connections before this module ever needed one. This lesson reuses that work directly: the same PostgresHook, the same cityflow_postgres Connection, now writing CityFlow’s real transformed summary instead of a demo table.
By the end of this lesson, you will be able to:
- Write a load task that creates its target table and inserts real rows via
PostgresHook.insert_rows() - Verify a load with a query whose result ties back to an earlier task’s own reported number, not just a standalone count
- Explain why load is the task most worth isolating from the rest of the pipeline
Reusing Ingest, Validate, and Transform Unchanged
cityflow_etl_load_demo_dag.py has four tasks. The first three — ingest_taxi_trips, validate_trips, transform_daily_borough_summary — are exactly Lessons 2 and 3’s real logic, unchanged, writing to their own etl_load_demo work directory. The only new task is the fourth:
# gate: skip
@task
def load_daily_borough_summary_to_postgres(transform_result: dict) -> int:
"""Write the real transformed summary to Postgres via PostgresHook
and the named 'cityflow_postgres' Connection -- no raw psycopg2, no
hardcoded host/port/password -- then prove it landed with a real
SELECT COUNT(*) and SELECT SUM(trip_count), the second one tying
back to transform's own reported total_trip_count."""
import pandas as pd
from airflow.providers.postgres.hooks.postgres import PostgresHook
summary = pd.read_csv(TRANSFORMED_PATH)
hook = PostgresHook(postgres_conn_id=POSTGRES_CONN_ID)
hook.run(f"DROP TABLE IF EXISTS {TABLE_NAME};")
hook.run(
f"""
CREATE TABLE {TABLE_NAME} (
id SERIAL PRIMARY KEY,
pickup_date DATE NOT NULL,
borough TEXT NOT NULL,
trip_count INTEGER NOT NULL,
avg_fare_amount REAL,
avg_trip_distance REAL,
avg_trip_duration_minutes REAL,
total_revenue REAL
);
"""
)
print(f"created table {TABLE_NAME} via PostgresHook.run()")
rows = list(
summary[
["pickup_date", "Borough", "trip_count", "avg_fare_amount", "avg_trip_distance",
"avg_trip_duration_minutes", "total_revenue"]
].itertuples(index=False, name=None)
)
hook.insert_rows(
table=TABLE_NAME,
rows=rows,
target_fields=["pickup_date", "borough", "trip_count", "avg_fare_amount",
"avg_trip_distance", "avg_trip_duration_minutes", "total_revenue"],
commit_every=50,
)
print(f"loaded {len(rows)} real summary rows into {TABLE_NAME} via PostgresHook.insert_rows()")
row_count = hook.get_first(f"SELECT COUNT(*) FROM {TABLE_NAME};")[0]
total_trip_count = hook.get_first(f"SELECT SUM(trip_count) FROM {TABLE_NAME};")[0]
print(f"SELECT COUNT(*) via hook confirms {row_count} rows landed in Postgres")
print(f"SELECT SUM(trip_count) via hook confirms {total_trip_count} real trips represented")
assert row_count == len(summary), f"expected {len(summary)} rows, got {row_count}"
assert total_trip_count == transform_result["total_trip_count"], (
f"loaded SUM(trip_count) {total_trip_count} != transform's "
f"total_trip_count {transform_result['total_trip_count']}"
)
return row_countThe verification inside the task itself is worth pausing on: it isn’t just SELECT COUNT(*) against the number of rows this task tried to insert (that would only prove the insert didn’t silently drop rows). SELECT SUM(trip_count) is checked against transform_result["total_trip_count"] — a number that traveled here through XCom all the way from transform_daily_borough_summary, which itself asserted that number against validate_trips’s rows_passed. Three tasks, three independent computations of essentially the same fact, all required to agree before this task returns successfully.
Triggered for Real: 7 Rows Land in Postgres
This lesson’s demo needs cityflow-lessons-db, the same throwaway Postgres container Module 4 introduced, and the same cityflow_postgres Connection — both already configured, nothing new to set up.
# gate: setup: docker rm -f cityflow-lessons-db >/dev/null 2>&1; docker run -d --name cityflow-lessons-db -e POSTGRES_USER=cityflow -e POSTGRES_PASSWORD=cityflow -e POSTGRES_DB=cityflow -p 55432:5432 postgres:13 >/dev/null && sleep 4
# gate: teardown: docker rm -f cityflow-lessons-db >/dev/null 2>&1
import json
import subprocess
import time
dag_id = "cityflow_etl_load_demo"
deadline = time.time() + 320
found = False
while time.time() < deadline:
result = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list", "--output", "json"],
capture_output=True, text=True, check=True,
)
if any(d["dag_id"] == dag_id for d in json.loads(result.stdout)):
found = True
break
time.sleep(5)
assert found, f"{dag_id} was never picked up by the scheduler"
subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
capture_output=True, text=True, check=True)
run_id = f"lesson4_m5_verify_{int(time.time())}"
trigger = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
capture_output=True, text=True,
)
if trigger.returncode != 0:
subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "reserialize"],
capture_output=True, text=True, check=True)
subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
capture_output=True, text=True, check=True,
)
print(f"triggered {dag_id!r} as {run_id!r}")
deadline = time.time() + 90
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
time.sleep(3)
r = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list-runs", "-d", dag_id, "--output", "json"],
capture_output=True, text=True, check=True,
)
runs = json.loads(r.stdout)
state = next(x["state"] for x in runs if x["run_id"] == run_id)
print(f"run {run_id!r} reached state: {state}")
assert state == "success", f"expected success, got {state}"
# Independent check: query cityflow-lessons-db directly, bypassing Airflow entirely.
import psycopg2
conn = psycopg2.connect(host="localhost", port=55432, user="cityflow", password="cityflow", dbname="cityflow")
cur = conn.cursor()
cur.execute("SELECT COUNT(*), SUM(trip_count) FROM cityflow_etl_load_demo;")
row_count, total_trips = cur.fetchone()
print(f"direct psycopg2 check: {row_count} rows, SUM(trip_count)={total_trips}")
cur.execute("SELECT pickup_date, borough, trip_count, avg_fare_amount FROM cityflow_etl_load_demo ORDER BY borough;")
for row in cur.fetchall():
print(" ", row)
assert row_count == 7, f"expected 7, got {row_count}"
assert total_trips == 69290, f"expected 69290, got {total_trips}"
conn.close()triggered 'cityflow_etl_load_demo' as 'lesson4_m5_verify_1784914602'
run 'lesson4_m5_verify_1784914602' reached state: success
direct psycopg2 check: 7 rows, SUM(trip_count)=69290
(datetime.date(2024, 1, 1), 'Bronx', 76, 26.79)
(datetime.date(2024, 1, 1), 'Brooklyn', 646, 27.45)
(datetime.date(2024, 1, 1), 'EWR', 20, 78.48)
(datetime.date(2024, 1, 1), 'Manhattan', 59000, 17.06)
(datetime.date(2024, 1, 1), 'Queens', 9174, 53.79)
(datetime.date(2024, 1, 1), 'Staten Island', 3, 96.4)
(datetime.date(2024, 1, 1), 'Unknown', 371, 40.06)The second check doesn’t go through Airflow, PostgresHook, or the DAG at all — it’s this script’s own psycopg2 connection straight to cityflow-lessons-db on localhost:55432. It agrees with the task’s own internal check, which agreed with transform_daily_borough_summary’s reported total, which agreed with validate_trips’s rows_passed: 69,290 real trips, counted the same way four separate times.
A Real Screenshot of All Four Real Stages
Key Takeaways
- Load is the one task that talks to an outside system, which is exactly why isolating it (Lesson 1’s argument) matters most here: a Postgres outage only ever fails this one task, not the transform work already sitting finished behind it.
PostgresHook.insert_rows(..., commit_every=50)batches real inserts the way a real loader would, the same pattern Module 4’s guided project used at a larger scale.- The load task’s own verification query checks against a number that traveled through three tasks’ worth of XCom, not a number this task invented for itself — that’s what makes “it loaded” mean something more than “some rows are in a table.”
- A load task’s
DROP TABLE IF EXISTS/CREATE TABLEpair (notCREATE TABLE IF NOT EXISTS) makes every real trigger idempotent — re-running this DAG against the same real data always produces the same real table, never a stale mix of old and new rows.
Exercises
Exercise 1 — Add a Postgres index on (pickup_date, borough) right after the CREATE TABLE call, and use hook.get_records("SELECT indexname FROM pg_indexes WHERE tablename = 'cityflow_etl_load_demo';") in a follow-up task to confirm the real index exists.
Exercise 2 — Change POSTGRES_CONN_ID to a connection ID that doesn’t exist and re-trigger. Read the real task failure in the log and confirm it’s a connection-lookup error, not a SQL error — the same distinction Module 4 Lesson 2’s exercises drew.
Exercise 3 — Add a fifth task, verify_load_row_count, that runs after load_daily_borough_summary_to_postgres with its own independent PostgresHook and its own SELECT COUNT(*), wired with trigger_rule="all_success" (the default) so it only runs if the load genuinely succeeded. Trigger the DAG and confirm this task’s own log shows the same count the load task already reported.
Previous lesson
Lesson 3 - Validate and Transform Tasks