Airflow Interview Questions for Data Engineers

import glob
import os
from datetime import datetime, timedelta, timezone

import pandas as pd
from airflow.exceptions import AirflowFailException
from airflow.sdk import dag, task

LANDING, WAREHOUSE = "landing", "warehouse"                     # one CSV per day arrives in landing/


def build(dag_id, fixed):
    def on_failure(context):                                     # the fixed version alerts from a callback
        open("alerts.log", "a").write(f"{dag_id} failed\n")

    @dag(dag_id=dag_id, schedule="@daily", start_date=datetime(2010, 3, 1, tzinfo=timezone.utc), catchup=False,
         default_args={"retries": 2, "retry_delay": timedelta(minutes=5)}, on_failure_callback=on_failure if fixed else None)
    def power_daily():
        @task
        def extract(logical_date=None, data_interval_start=None, data_interval_end=None):
            day = logical_date.strftime("%Y-%m-%d")              # the date this run is for
            rows = pd.read_csv(f"{LANDING}/{day}.csv")
            with open(f"interval_{dag_id}_{day}.txt", "w") as f:
                f.write(f"{data_interval_start:%Y-%m-%d %H:%M} -> {data_interval_end:%Y-%m-%d %H:%M}")
            return {"day": day, "rows": len(rows), "missing": int(rows.Global_active_power.isna().sum())}

        @task
        def validate(batch):
            if batch["missing"] > 0.1 * batch["rows"]:           # retrying will not bring readings back
                raise AirflowFailException(f"{batch['day']}: {batch['missing']} of {batch['rows']} readings missing")
            return batch

        @task
        def load(batch):
            out = f"{WAREHOUSE}/day={batch['day']}"
            os.makedirs(out, exist_ok=True)
            pd.read_csv(f"{LANDING}/{batch['day']}.csv").to_csv(f"{out}/part.csv", index=False)   # overwrite: reruns are safe
            return batch["day"]

        @task
        def meter_total(day, meter):                             # one mapped task instance per sub-meter
            return float(pd.read_csv(f"{WAREHOUSE}/day={day}/part.csv")[f"Sub_metering_{meter}"].sum())

        @task(trigger_rule="all_done")
        def cleanup():
            for f in glob.glob("tmp_*"):
                os.remove(f)

        checked = validate(extract())
        totals = meter_total.partial(day=load(checked)).expand(meter=[1, 2, 3])
        if fixed:
            totals >> cleanup().as_teardown()                    # teardowns do not decide the run's state
        else:
            @task(trigger_rule="one_failed")
            def alert():
                open("alerts.log", "a").write(f"{dag_id} alert task ran\n")
            [checked, totals] >> alert()
            totals >> cleanup()

    return power_daily()


naive, fixed = build("power_daily", fixed=False), build("power_daily_fixed", fixed=True)

if __name__ == "__main__":
    raw = pd.read_csv("household_power_consumption.txt", sep=";", na_values="?")   # UCI household power
    raw["day"] = pd.to_datetime(raw.Date, format="%d/%m/%Y").dt.strftime("%Y-%m-%d")
    os.makedirs(LANDING, exist_ok=True)
    for d in ("2010-03-19", "2010-03-20"):
        raw[raw.day == d].drop(columns="day").to_csv(f"{LANDING}/{d}.csv", index=False)
    for pipeline, d in ((naive, "2010-03-19"), (naive, "2010-03-20"), (naive, "2010-03-19"), (fixed, "2010-03-20")):
        run = pipeline.test(logical_date=datetime.fromisoformat(d).replace(tzinfo=timezone.utc))
        states = {(ti.task_id if ti.map_index < 0 else f"{ti.task_id}[{ti.map_index}]"): ti.state for ti in run.get_task_instances()}
        rows = sum(len(pd.read_csv(f)) for f in glob.glob(f"{WAREHOUSE}/day=*/part.csv"))
        print(f"{pipeline.dag_id} {d}: run {run.state} | data interval {open(f'interval_{pipeline.dag_id}_{d}.txt').read()} | warehouse rows {rows:,}")
        print("   ", ", ".join(f"{k}={v}" for k, v in sorted(states.items())))
    print("alerts.log:", open("alerts.log").read().strip().replace("\n", " | "))

power_daily 2010-03-19: run success | data interval 2010-03-19 00:00 -> 2010-03-19 00:00 | warehouse rows 1,440
    alert=skipped, cleanup=success, extract=success, load=success, meter_total[0]=success, meter_total[1]=success, meter_total[2]=success, validate=success
power_daily 2010-03-20: run success | data interval 2010-03-20 00:00 -> 2010-03-20 00:00 | warehouse rows 1,440
    alert=success, cleanup=success, extract=success, load=upstream_failed, meter_total[0]=upstream_failed, meter_total[1]=upstream_failed, meter_total[2]=upstream_failed, validate=failed
power_daily 2010-03-19: run success | data interval 2010-03-19 00:00 -> 2010-03-19 00:00 | warehouse rows 1,440
    alert=skipped, cleanup=success, extract=success, load=success, meter_total[0]=success, meter_total[1]=success, meter_total[2]=success, validate=success
power_daily_fixed 2010-03-20: run failed | data interval 2010-03-20 00:00 -> 2010-03-20 00:00 | warehouse rows 1,440
    cleanup=success, extract=success, load=upstream_failed, meter_total[0]=upstream_failed, meter_total[1]=upstream_failed, meter_total[2]=upstream_failed, validate=failed
alerts.log: power_daily alert task ran | power_daily_fixed failed