# Run Rivet on Apache Airflow

Rivet extracts your tables; Airflow schedules and watches them. This recipe makes
Airflow's graph **be** Rivet's extraction plan — small tables parallelised, heavy
tables run one at a time, a barrier between waves — with per-table retries, logs,
and alerting for free, and a real `rivet` binary doing the work. It builds **one
DAG per source database** (PostgreSQL, MySQL, SQL Server) from a single factory.

```
docs/recipes/airflow/
├── Dockerfile                 # apache/airflow + the rivet release binary baked in
├── docker-compose.2.10.yaml   # official Airflow 2.10 stack, adapted
└── dags/
    ├── rivet_waves_dag.py      # factory → rivet_waves_{postgres,mysql,mssql}
    ├── postgres.yaml / mysql.yaml / mssql.yaml      # the configs you edit  ← yours
    └── postgres.plan.json / mysql.plan.json / …     # rivet plan output (auto-refreshed)
```

---

## Try it locally

This is the **official** Apache Airflow docker-compose stack (Postgres metadata +
Redis + CeleryExecutor — not SQLite/Sequential), adapted three ways: example DAGs
are off, the worker image has `rivet` baked in, and a dedicated TLS-enabled
Postgres holds Rivet's durable run state.

```bash
cd docs/recipes/airflow
docker compose -f docker-compose.2.10.yaml up --build      # Airflow 2.10.5
```

First boot builds the image and migrates the metadata DB (~2-3 min). Then open
<http://localhost:8080> (login **`airflow`** / **`airflow`**). There are three DAGs
— **`rivet_waves_postgres`**, **`rivet_waves_mysql`**, **`rivet_waves_mssql`** —
one per source. Un-pause and trigger one, and open **Graph**:

```
plan ──> wave_2 ──────────> wave_3 ───────────────> wave_4
         [small tables       [bench_decimal →        [bench_narrow →
          in parallel]        bench_hc →              content_items]
                              orders → bench_wide]    (one at a time)
```

The demo runs against the project's fixtures on the host (Postgres :5432, MySQL
:3306, SQL Server :1433), so `docker compose up` in the rivet repo first.
`docker compose -f … down -v` tears everything down (`-v` drops the state too).

(The DAG keeps a `try/except` around the `BashOperator` import so it still parses
on Airflow 3.x — where it moved to the `standard` provider — if you point it at a
3.x stack yourself.)

---

## What the graph does

**`plan` → waves.** `rivet plan` scores every export (size, cursor quality, chunk
geometry, risk) and groups them into waves. The DAG runs the waves lowest-first
with a **barrier** between them — wave N+1 starts only after every task in wave N
succeeds.

**Cheap parallel, heavy serial — within a wave.** This is the part that matters:
only the cheap exports (planner `cost_class: low`) run in parallel. The heavier
ones run **one at a time** (a sequential chain). The whole reason the planner
defers big tables into a late wave is to *not* pile several large scans onto the
source at once — so running them in parallel would defeat the point. (Same split
as the in-engine `rivet apply --parallel-export-processes`.)

**Each task is a real `rivet run`.** A wave task is `rivet run --config
<source>.yaml --export <table>` — a single table, with `--reconcile` (source
`COUNT(*)` vs exported rows) on a fresh full/chunked export. The task log is the
actual rivet output:

```
✓ orders         incremental  250,000 rows  1 files  3.4 MB  2.7s  RSS 64 MB
```

A failing reconcile (or any error) exits non-zero with a stable `[RIVET_*]` code,
so the task fails loudly and the wave barrier stops everything downstream — before
a half-extracted table reaches a warehouse load.

---

## Config → plan → graph (no manual steps)

The graph is generated from Rivet's planner, and each DAG's `plan` task keeps it
fresh:

1. **`dags/<source>.yaml`** is the config you edit (`postgres.yaml` etc).
   `rivet init --source <url>` scaffolds one from your live schema (use
   `--exclude '<glob>'` to drop test / junk tables); the checked-in samples are
   the fixtures trimmed to a clean set.
2. The DAG's first task, **`plan`**, runs `rivet plan --format json` and
   **atomically** rewrites `<source>.plan.json` (temp file, swapped in only if
   rivet succeeded *and* the output is valid JSON — a failed plan can't truncate
   the graph source and break the DAG).
3. The DAG reads `<source>.plan.json` at **parse time** to lay out the waves. So
   editing the config and re-running the DAG re-shapes the graph on the next parse
   — no hand-run CLI, no committing the plan by yourself. The checked-in sample
   makes the DAG work on the very first boot.

Same-named tables across engines never collide: each source writes to its own
`./output/<engine>/<table>/` and keeps its own state database, so a `bench_hc` in
Postgres and one in MySQL don't share cursor / shape / file-log state.

### Skip tables — you don't have to extract everything

Set an env var on the workers (or in the compose `environment:`):

| Variable        | Effect                                                            |
|-----------------|------------------------------------------------------------------|
| `RIVET_EXCLUDE` | Comma-separated tables to drop. A wave left empty disappears.     |
| `RIVET_ONLY`    | Comma-separated allow-list — run only these.                     |

---

## Many sources → one bucket

For a shared S3 / GCS data lake, point each source's config at the **same bucket
with a per-source prefix** so same-named tables across engines never collide:

```yaml
# postgres.s3.yaml
destination:
  type: s3
  bucket: my-data-lake                 # ← one bucket for every source
  prefix: rivet/postgres/{export}/     # ← namespaced by source; {export} = table
# mysql.s3.yaml  → prefix: rivet/mysql/{export}/
# mssql.s3.yaml  → prefix: rivet/mssql/{export}/
```

A `bench_hc` then lands at `s3://my-data-lake/rivet/postgres/bench_hc/`, the MySQL
one at `…/rivet/mysql/bench_hc/` — separate objects. The namespace is **required**,
not cosmetic: different databases are different data, and the *same* logical table
yields different Arrow types per engine (`uuid_col` is `FixedSizeBinary(16)` on
PG/MSSQL but `Utf8` on MySQL; MySQL drops the timestamp's UTC tz) — merging them
under one prefix would write a dataset with incompatible schemas.

The checked-in `*.s3.yaml` / `postgres.gcs.yaml` configs build the cloud DAGs
(`rivet_waves_<source>_s3`, `rivet_waves_postgres_gcs`); the demo writes to the
project's MinIO / fake-gcs emulators (worker env supplies `AWS_ACCESS_KEY_ID` /
`AWS_SECRET_ACCESS_KEY`). A real bucket drops the `endpoint` + `allow_anonymous`
and uses normal cloud credentials.

---

## Recovery — all from Airflow, no shell

A chunked export checkpoints per chunk, so a kill mid-chunk is recoverable. The
whole recovery story is operable from the Airflow UI:

| Situation | What to do |
|-----------|------------|
| **Transient crash mid-chunk** (worker killed, timeout, retry) | **Nothing — automatic.** The task tries a clean run; if rivet reports the checkpoint is still in progress, it `--resume`s from the last good chunk. (Keyed on the checkpoint *state*, not the retry count — idempotent: a fresh run, a mid-chunk crash, and a finished export each do the right thing.) |
| **Unresumable checkpoint** (chunk params changed, or you want a clean re-extract) | **Admin → Variables → `rivet_reset`** = comma-list of tables → trigger the DAG. Each listed table's checkpoint is wiped (`rivet state reset-chunks`) before its run. Clear the Variable afterwards. |
| **Anything else** | the normal task **Clear / re-trigger** in the UI. |

A resumed run writes only the chunks left after the crash, so the reconcile gate
is dropped on the resume path (a full-table `COUNT(*)` would false-mismatch the
remainder) — the same reasoning as skipping reconcile on incremental exports.

---

## Architecture

- **`rivet` binary** — baked into the worker image from the published release
  ([`Dockerfile`](Dockerfile), arch-aware: pulls the `x86_64` or `aarch64` linux
  build). To pin a version, set `RIVET_VERSION` in the compose `build.args`.
- **Durable state** — `rivet-state-db`, a dedicated TLS Postgres service, with a
  **separate database per source** (`rivet_state_postgres` / `_mysql` / `_mssql`);
  the factory passes each DAG its own `RIVET_STATE_URL` with `sslmode=require`.
  SQLite state in a bind-mount corrupts under parallel writers; Rivet refuses to
  send state credentials in cleartext to a non-loopback host (CWE-319), so the
  service speaks TLS (a self-signed cert generated at startup — fine in-cluster).
- **Source credentials — Airflow Connections.** Each source is an Airflow
  Connection (`rivet_postgres` / `rivet_mysql` / `rivet_mssql`), editable in the UI
  (Admin → Connections); the demo seeds them via `AIRFLOW_CONN_*` in the compose so
  it works out of the box. The DAG builds the URL from the connection's **fields**
  in rivet's scheme — not `get_uri()`, whose scheme differs per engine (Airflow
  emits `mssql://`, rivet wants `sqlserver://`) — and exports it as the env var the
  config's `url_env` reads. Nothing about a database is hard-coded in the DAG. The
  demo opts into plaintext to the fixtures with `source.tls: { mode: disable }`; a
  real remote DB uses `mode: verify-full`.
- **State schema is created up front.** `initdb` creates the state *databases*, but
  rivet's *tables* are created on first connect — under parallelism that races
  ("create version table: db error"). A one-shot `rivet-state-schema-init` service
  runs `rivet state show` against every state DB at boot, so the schema exists
  before any wave task runs. The Airflow services wait for it.

---

## Why this shape

Rivet's planner is deliberately **advisory** ([ADR-0006](../../adr/0006-source-aware-prioritization.md)):
it scores and groups, it does not schedule. That's what lets a real scheduler own
execution — retries, backfills, SLAs, alerting — while Rivet owns *source safety*
(which tables to defer, which to run serially, how hard to push the database).
This recipe is the seam between the two: the planner's `waves` and cost classes
become Airflow's graph, and nothing about your database or its credentials leaves
your environment.
