Solutions — Homework 06

Lesson 06 · Data Engineering with Python

Task 1 — Make the daily rollup idempotent

Import the helpers at the top of gharchive_daily_rollup.py and refactor _rollup:

import pathlib
import sys

# Add the dags directory to the path so _helpers is importable
sys.path.insert(0, str(pathlib.Path(__file__).parent))
from _helpers import atomic_replace, transactional_write_path

def _rollup(**context):
    import glob
    import pandas as pd

    day = context["data_interval_start"].strftime("%Y-%m-%d")
    files = sorted(glob.glob(
        f"/opt/airflow/sample-pipeline/data/raw/{day}-*.json.gz"
    ))
    if not files:
        print(f"no raw files for {day}, skipping")
        return

    # Read and aggregate (implementation depends on Lesson 01 transform)
    dfs = [pd.read_json(f, lines=True, compression="gzip") for f in files]
    df = pd.concat(dfs, ignore_index=True)
    out = df.groupby("type").size().reset_index(name="count")

    final = pathlib.Path(
        f"/opt/airflow/sample-pipeline/data/daily_rollups/{day}.parquet"
    )
    final.parent.mkdir(parents=True, exist_ok=True)
    staging = transactional_write_path(final, context["run_id"])
    out.to_parquet(staging, index=False)
    atomic_replace(staging, final)

Common mistake: calling mkdir on the staging path's parent instead of the final path's parent. The staging file lives next to the final file (same directory), so final.parent.mkdir() covers both. Creating a separate staging directory defeats the purpose of the pattern — the atomic rename must cross no filesystem boundary.

Task 2 — Tune retry budgets with a failure-notification task

from airflow.operators.bash import BashOperator

notify_slack = BashOperator(
    task_id="notify_slack",
    bash_command='echo "would notify slack: the hour {{ ds }} transform failed"',
    trigger_rule="one_failed",
    retries=0,
)
transform >> notify_slack

Common mistake 1: placing notify_slack upstream of transform (notify_slack >> transform). The notification needs to fire after failure, not before. The correct direction is transform >> notify_slack.

Common mistake 2: using trigger_rule="all_failed". This fires only when every upstream task has failed. Since extract is still upstream and may have succeeded, all_failed would never trigger here. Use trigger_rule="one_failed" to fire when at least one upstream task has failed.

Task 3 — Identify the non-idempotent operation

(a) — Not idempotent.

def task(**ctx):
    with open("counter.txt", "a") as f:
        f.write("1\n")

Append mode adds a line on every run. After two runs there are two lines; after three runs, three. Fix: write the content as the full file using overwrite mode (mode='w'), or use a staging-then-rename to make the overwrite atomic.

(b) — Not idempotent.

def task(**ctx):
    db.execute(
        "INSERT INTO events (hour, count) VALUES (%s, %s)",
        (hour, n),
    )

A plain INSERT raises a primary-key (or unique-key) violation on the second run. Airflow retries the task, masking the real problem. Fix: write an upsert:

INSERT INTO events (hour, count)
VALUES (%s, %s)
ON CONFLICT (hour) DO UPDATE SET count = EXCLUDED.count;

(c) — Not idempotent.

from uuid import uuid4

def task(**ctx):
    requests.post(
        "https://api.example.com/orders",
        json={"id": uuid4().hex, "amount": 100},
    )

A new UUID is generated on every call, so every retry creates a new order. The receiving API has no way to detect the duplicate. Fix: derive the order ID from a stable run key (e.g., context["run_id"] or a hash of data_interval_start) and pass it as an Idempotency-Key header or as the explicit ID. The API can then reject or de-duplicate duplicate requests for the same key.

(d) — Idempotent.

def task(**ctx):
    df.to_parquet(f"out/{hour}.parquet")  # pandas default: overwrite

Pandas writes the complete file before returning. Calling it a second time with the same data produces the same bytes at the same path. The output is deterministic.

Caveat: while logically idempotent, this write is not atomic. A crash mid-write leaves a partial file at the final path. For production use, combine the overwrite with staging-then-rename to get both idempotency and atomicity.

Module 3 introduces dbt, which makes most of the idempotency work disappear behind table materializations — dbt's incremental models generate the upsert SQL automatically.