/airflow-state-store

name: airflow-state-store
description: Persists task and asset state across retries and DAG runs using Airflow 3.3's AIP-103 key/value stores (`task_state_store`, `asset_state_store`) and the crash-safe `ResumableJobMixin`. Use when the user asks about task state store, checkpointing in tasks, persisting state across retries, job IDs surviving worker crashes, watermarks, asset metadata, resumable tasks, crash-safe operators, or "what's new in Airflow 3.3". Also use proactively when reading a DAG that uses Variables or XCom for intra-task coordination state (flag the anti-pattern), or when reviewing any DAG that submits a job to an external system (Databricks, Snowflake, BigQuery, Redshift, Spark, dbt Cloud, EMR, AWS Batch, etc.) and waits for it to finish, whether as one submit-and-wait operator or split into submit + sensor/polling tasks; covers `wait_for_termination`, `deferrable`, `durable`, and hand-rolled polling sensors. State-persistence pieces require Airflow 3.3+; submit+poll architecture guidance applies on any version.

Airflow Task State Store (AIP-103)

Airflow 3.3 ships two key/value stores and a crash-safety mixin for operators that submit external jobs.

> **`task_state_store`, `asset_state_store`, and `ResumableJobMixin`'s crash-safety guarantee require Airflow 3.3+.** Check first: > ```bash > af config version > ``` > Below 3.3: `task_state_store`/`asset_state_store` are unavailable, and `durable=True` is a no-op — provider operators ship a pre-3.3 `ResumableJobMixin` shim that always submits fresh (see Section 5). Tell the user those specific features aren't available yet and link the AIP-103 tracking issue. This does **not** gate Section 6's Triggerer-vs-`mode="reschedule"` decision, or the general "green submit ≠ success" anti-pattern — those apply on any Airflow version. On a pre-3.3 DAG, give that guidance in full; only drop the "`durable=True` adds crash-safety" half of it.

---

Section 1 — Pick the right primitive

| I need to… | Use | |---|---| | Persist a cursor, offset, or job ID so a retry can resume instead of restart | `task_state_store` | | Pass small coordination state within one task across retries (not between tasks) | `task_state_store` | | Store a watermark or last-processed timestamp per asset, surviving across DAG runs | `asset_state_store` | | Cache asset-level metadata (manifest hash, row count, schema version) | `asset_state_store` | | Make an existing non deferrable operator crash-safe when it submits to an external system | `task_state_store` or `ResumableJobMixin` |

**When NOT to use these:**

• Passing data *between* tasks -> use XCom
• Large payloads (model weights, dataframes) -> use XCom with an object storage backend
• Config or secrets shared across DAGs -> use Variables or Connections

---

Section 2 — Detect anti-patterns in existing DAGs (on demand)

When the user asks to review a DAG or asks "is there a better way", scan for these patterns and flag them:

| Pattern seen in DAG | Problem | Recommend | |---|---|---| | `Variable.get(...)` / `Variable.set(...)` inside a `@task` body for per-run state | Variables are global and shared; no scoping to task instance or retry | `task_state_store` | | `context["ti"].xcom_push(key="job_id", ...)` to survive retries | XCom is scoped to a DAG run, not a retry; a new ti_id is issued per retry | `task_state_store` or `ResumableJobMixin` | | Manual `if Variable.get("job_id"): reconnect else: submit` retry-resume logic | Reimplements what `ResumableJobMixin` already provides, without the crash-safety guarantee | `ResumableJobMixin` | | `Variable.set("last_processed_at", ...)` for watermarks | Global; any DAG or task can overwrite it; no scoping to asset | `asset_state_store` | | Separate `submit` task (`wait_for_termination=False` / fire-and-forget) + a second sensor/polling task waiting on the same external job (Databricks, Snowflake, BigQuery, Redshift, Spark, etc.) | A green submit task only means the job was *accepted*, not that it *succeeded* — only the sensor task's outcome reflects reality. | See **Section 6, "Submit-and-poll DAGs: one task or two?"** — the right call depends on Triggerer availability and job duration, not a single fixed answer. |

Show a before/after snippet when flagging. Use the canonical examples in Steps 3–5 as the "after".

**The submit+sensor split is worth a comment even when the sensor code is bug-free** — reviewing the sensor's code quality (correct `mode=`, correct terminal-state handling, cached hook) is a separate question from whether the two-task split is the right architecture. Follow **Section 6, "Submit-and-poll DAGs: one task or two?"** for that decision; don't re-derive it here.

---

Section 3 — `task_state_store`: per-task coordination state

`task_state_store` is a key/value store scoped to a single task instance identity (dag_id + run_id + task_id + map_index). It survives retries — a new retry on the same task reads the same store.

from airflow.sdk import dag, task
from pendulum import datetime

@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
def etl_with_checkpoint():

    @task(retries=3)
    def process_records(**context):
        task_state_store = context["task_state_store"]  # injected by Airflow, no setup needed
        cursor = task_state_store.get("last_cursor", default=0)
        records = fetch_records_after(cursor)
        for record in records:
            process(record)
            cursor = record["id"]
            task_state_store.set("last_cursor", cursor)   # checkpoint after each record

    process_records()

etl_with_checkpoint()

**API:**

from airflow.sdk import NEVER_EXPIRE

task_state_store.get(key, default=None)                        # returns a JsonValue or default
task_state_store.set(key, value)                               # uses default_retention_days
task_state_store.set(key, value, retention=timedelta(days=7))  # per-key TTL override
task_state_store.set(key, value,