airflow-state-store

Name: airflow-state-store
Author: astronomer

33installs55forks

How to install airflow-state-store?

astronomer/agents/airflow-state-store

$npx -y skills add astronomer/agents --skill airflow-state-store

Installs into the current project.

Summary

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

Files · 1

View on GitHub

SKILL.md

1	# Airflow Task State Store (AIP-103)
2
3	Airflow 3.3 ships two key/value stores and a crash-safety mixin for operators that submit external jobs.
4
5	> Requires Airflow 3.3+. Check first:
6	> ```bash
7	> af config version
8	> ```
9	> If the version is below 3.3, tell the user these features are not yet available and link them to the AIP-103 tracking issue instead.
10
11	---
12
13	## Section 1 — Pick the right primitive
14
15	\| I need to… \| Use \|
16	\|---\|---\|
17	\| Persist a cursor, offset, or job ID so a retry can resume instead of restart \| `task_state_store` \|
18	\| Pass small coordination state within one task across retries (not between tasks) \| `task_state_store` \|
19	\| Store a watermark or last-processed timestamp per asset, surviving across DAG runs \| `asset_state_store` \|
20	\| Cache asset-level metadata (manifest hash, row count, schema version) \| `asset_state_store` \|
21	\| Make an existing non deferrable operator crash-safe when it submits to an external system \| `task_state_store` or `ResumableJobMixin` \|
22
23	When NOT to use these:
24	- Passing data between tasks -> use XCom
25	- Large payloads (model weights, dataframes) -> use XCom with an object storage backend
26	- Config or secrets shared across DAGs -> use Variables or Connections
27
28	---
29
30	## Section 2 — Detect anti-patterns in existing DAGs (on demand)
31
32	When the user asks to review a DAG or asks "is there a better way", scan for these patterns and flag them:
33
34	\| Pattern seen in DAG \| Problem \| Recommend \|
35	\|---\|---\|---\|
36	\| `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` \|
37	\| `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` \|
38	\| Manual `if Variable.get("job_id"): reconnect else: submit` retry-resume logic \| Reimplements what `ResumableJobMixin` already provides, without the crash-safety guarantee \| `ResumableJobMixin` \|
39	\| `Variable.set("last_processed_at", ...)` for watermarks \| Global; any DAG or task can overwrite it; no scoping to asset \| `asset_state_store` \|
40
41	Show a before/after snippet when flagging. Use the canonical examples in Steps 3–5 as the "after".
42
43	---
44
45	## Section 3 — `task_state_store`: per-task coordination state
46
47	`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.
48
49	```python
50	from airflow.sdk import dag, task
51	from pendulum import datetime
52
53	@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
54	def etl_with_checkpoint():
55
56	@task(retries=3)
57	def process_records(**context):
58	task_state_store = context["task_state_store"] # injected by Airflow, no setup needed
59	cursor = task_state_store.get("last_cursor", default=0)
60	records = fetch_records_after(cursor)
61	for record in records:
62	process(record)
63	cursor = record["id"]
64	task_state_store.set("last_cursor", cursor) # checkpoint after each record
65
66	process_records()
67
68	etl_with_checkpoint()
69	```
70
71	API:
72	```python
73	from airflow.sdk import NEVER_EXPIRE
74
75	task_state_store.get(key, default=None) # returns a JsonValue or default
76	task_state_store.set(key, value) # uses default_retention_days
77	task_state_store.set(key, value, retention=timedelta(days=7)) # per-key TTL override
78	task_state_store.set(key, value, retention=NEVER_EXPIRE) # never expires regardless of config
79	task_state_store.delete(key) # no-op if key does not exist
80	task_state_store.clear() # delete all keys for this task instance
81	```
82
83	Key rules:
84	- Values must be JSON-serializable (`str`, `int`, `float`, `bool`, `list`, `dict` — `None` values are rejected).
85	- Default expiry is controlled by `[state_store] default_retention_days` (0 = never expire).
86	- Use `NEVER_EXPIRE` for keys that must outlive the default retention window (e.g. a job ID for a multi-day Spark job).
87	- 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_

People who liked this also viewed

6 picks

Type

airflow-state-store

33installs55forks

View on GitHub

How to install airflow-state-store?

astronomer/agents/airflow-state-store

$npx -y skills add astronomer/agents --skill airflow-state-store

Installs into the current project.

Summary

Files · 1

View on GitHub

SKILL.md

1	# Airflow Task State Store (AIP-103)
2
3	Airflow 3.3 ships two key/value stores and a crash-safety mixin for operators that submit external jobs.
4
5	> Requires Airflow 3.3+. Check first:
6	> ```bash
7	> af config version
8	> ```
9	> If the version is below 3.3, tell the user these features are not yet available and link them to the AIP-103 tracking issue instead.
10
11	---
12
13	## Section 1 — Pick the right primitive
14
15	\| I need to… \| Use \|
16	\|---\|---\|
17	\| Persist a cursor, offset, or job ID so a retry can resume instead of restart \| `task_state_store` \|
18	\| Pass small coordination state within one task across retries (not between tasks) \| `task_state_store` \|
19	\| Store a watermark or last-processed timestamp per asset, surviving across DAG runs \| `asset_state_store` \|
20	\| Cache asset-level metadata (manifest hash, row count, schema version) \| `asset_state_store` \|
21	\| Make an existing non deferrable operator crash-safe when it submits to an external system \| `task_state_store` or `ResumableJobMixin` \|
22
23	When NOT to use these:
24	- Passing data between tasks -> use XCom
25	- Large payloads (model weights, dataframes) -> use XCom with an object storage backend
26	- Config or secrets shared across DAGs -> use Variables or Connections
27
28	---
29
30	## Section 2 — Detect anti-patterns in existing DAGs (on demand)
31
32	When the user asks to review a DAG or asks "is there a better way", scan for these patterns and flag them:
33
34	\| Pattern seen in DAG \| Problem \| Recommend \|
35	\|---\|---\|---\|
36	\| `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` \|
37	\| `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` \|
38	\| Manual `if Variable.get("job_id"): reconnect else: submit` retry-resume logic \| Reimplements what `ResumableJobMixin` already provides, without the crash-safety guarantee \| `ResumableJobMixin` \|
39	\| `Variable.set("last_processed_at", ...)` for watermarks \| Global; any DAG or task can overwrite it; no scoping to asset \| `asset_state_store` \|
40
41	Show a before/after snippet when flagging. Use the canonical examples in Steps 3–5 as the "after".
42
43	---
44
45	## Section 3 — `task_state_store`: per-task coordination state
46
47	`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.
48
49	```python
50	from airflow.sdk import dag, task
51	from pendulum import datetime
52
53	@dag(start_date=datetime(2025, 1, 1), schedule="@daily")
54	def etl_with_checkpoint():
55
56	@task(retries=3)
57	def process_records(**context):
58	task_state_store = context["task_state_store"] # injected by Airflow, no setup needed
59	cursor = task_state_store.get("last_cursor", default=0)
60	records = fetch_records_after(cursor)
61	for record in records:
62	process(record)
63	cursor = record["id"]
64	task_state_store.set("last_cursor", cursor) # checkpoint after each record
65
66	process_records()
67
68	etl_with_checkpoint()
69	```
70
71	API:
72	```python
73	from airflow.sdk import NEVER_EXPIRE
74
75	task_state_store.get(key, default=None) # returns a JsonValue or default
76	task_state_store.set(key, value) # uses default_retention_days
77	task_state_store.set(key, value, retention=timedelta(days=7)) # per-key TTL override
78	task_state_store.set(key, value, retention=NEVER_EXPIRE) # never expires regardless of config
79	task_state_store.delete(key) # no-op if key does not exist
80	task_state_store.clear() # delete all keys for this task instance
81	```
82
83	Key rules:
84	- Values must be JSON-serializable (`str`, `int`, `float`, `bool`, `list`, `dict` — `None` values are rejected).
85	- Default expiry is controlled by `[state_store] default_retention_days` (0 = never expire).
86	- Use `NEVER_EXPIRE` for keys that must outlive the default retention window (e.g. a job ID for a multi-day Spark job).
87	- 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_

People who liked this also viewed

6 picks

Type