astronomer/agents/airflow-state-store

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

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:

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, retention=NEVER_EXPIRE)       # never expires regardless of config
task_state_store.delete(key)                                   # no-op if key does not exist
task_state_store.clear()                                       # delete all keys for this task instance

Key rules:

• Values must be JSON-serializable (str, int, float, bool, list, dict — None values are rejected).
• Default expiry is controlled by [state_store] default_retention_days (0 = never expire).
• Use NEVER_EXPIRE for keys that must outlive the default retention window (e.g. a job ID for a multi-day Spark job).
• Max value size defaults to 64 KB; configurable via [state_store] max_value_storage_bytes (0 = no limit). For larger payloads, configure a custom [state_store] backend or a worker side backend configured via: [workers] state_store_backend.

Mapped tasks — each index has its own namespace:

When a task is dynamically mapped (task.expand(...)), each map index gets an isolated task_state_store scoped to its own map_index. Indices do not share state.

@task(retries=2)
def process_partition(partition_id, **context):
    task_state_store = context["task_state_store"]
    # Scoped to THIS index only — other indices have their own copy
    cursor = task_state_store.get("cursor", default=0)
    task_state_store.set("cursor", new_cursor)
process_partition.expand(partition_id=[0, 1, 2, 3])

clear() clears only the current index. To wipe state across all map indices of a task group, use the CLI or core API.

Before (anti-pattern):

@task
def process(**context):
    cursor = Variable.get("etl_cursor", default_var=0)
    # ... process ...
    Variable.set("etl_cursor", new_cursor)  # global, any task can overwrite

After:

@task(retries=3)
def process(**context):
    task_state_store = context["task_state_store"]
    cursor = task_state_store.get("cursor", default=0)
    # ... process ...
    task_state_store.set("cursor", new_cursor)    # scoped to this task instance

Section 4 — asset_state_store: per-asset metadata across DAG runs

asset_state_store is scoped to an asset, not a task instance. It persists across DAG runs — the same key on the same asset is readable and writable by any task that produces or consumes it.

from airflow.sdk import DAG, Asset, task
from datetime import datetime, timezone
ORDERS = Asset(name="orders/daily", uri="s3://warehouse/orders/daily")
with DAG(dag_id="producer", schedule=None, start_date=datetime(2026, 1, 1), catchup=False):
    @task(inlets=[ORDERS], outlets=[ORDERS])
    def load(asset_state_store=None):        # asset_state_store injected by Airflow — declare as a kwarg
        asset_state_store = asset_state_store[ORDERS]
        watermark = asset_state_store.get("watermark", default="2026-01-01T00:00:00+00:00")
        records = fetch_records_since(watermark)
        now = datetime.now(tz=timezone.utc).isoformat()
        asset_state_store.set("watermark", now)
        asset_state_store.set("last_run_summary", {"rows_loaded": len(records), "completed_at": now})
    load()

Reading the store from a consumer DAG:

with DAG(dag_id="consumer", schedule=[ORDERS], start_date=datetime(2026, 1, 1), catchup=False):
    @task(inlets=[ORDERS])
    def consume(asset_state_store=None):
        asset_state_store = asset_state_store[ORDERS]
        summary = asset_state_store.get("last_run_summary") or {}
        print(f"Processing {summary.get('rows_loaded')} rows up to {asset_state_store.get('watermark')}")
    consume()

Key rules:

• asset_state_store is injected by Airflow as a named kwarg — declare it as def my_task(asset_state_store=None). Do NOT combine with **context; Airflow injects it separately.
• Use datetime.now(tz=timezone.utc).isoformat() for timestamps — never datetime.utcnow() (not timezone-aware).
• Same JSON-serializable value constraint as task_state_store.
• No per-key expiry — asset state store entries have no TTL (the asset outlives any single run).
• Readable by any DAG that declares the asset as an inlet or outlet.

Mapped tasks — last writer wins:

asset_state_store is scoped to the asset, not the map index. If multiple mapped indices write the same key concurrently, the last write wins. Use distinct keys per index or ensure only one index writes to a given key.

@task(outlets=[my_asset])
def load_partition(partition_id, asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    # Distinct key per index — no race condition
    asset_state_store.set(f"offset_{partition_id}", new_offset)

Before (anti-pattern):

Variable.set(f"watermark_{asset_name}", new_offset)   # global, not scoped to asset

After:

@task(inlets=[my_asset], outlets=[my_asset])
def load(asset_state_store=None):
    asset_state_store = asset_state_store[my_asset]
    asset_state_store.set("watermark", new_offset)

Section 5 — ResumableJobMixin: crash-safe external job submission

Use whenever a task submits a job to an external system (Spark, Databricks, dbt Cloud, AWS Batch, etc.) and could resubmit a duplicate on retry — whether that same task also polls for completion inside one execute() call, or hands the job id to a separate downstream poll/sensor task. Without this mixin (or a provider operator that already builds it in, like DatabricksSubmitRunOperator's durable=True default), a worker crash after submission means the next retry of the submit task resubmits a duplicate job — that risk exists regardless of whether polling happens in that same task or a separate one.

Scope check before recommending it: if submit and poll are already two separate tasks (e.g. a submit task handing a job id to a downstream sensor/poll task via XCom), the poll/sensor task doesn't need this mixin — it never submits anything, so retrying it is already safe. The submit task still does; splitting off the poll step doesn't make the submit task's own duplicate-submission risk go away. See the table row below.

When NOT to use `ResumableJobMixin`: