AIP-103: Setting up foundation with ORM and Backend Abstraction

[amoghrajesh]

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

AIP-103 PR 1: DB Schema + Backend Abstraction

This PR lays the foundation for AIP-103 (Task State Management) — a key-value state store scoped to task instances and assets. It implements the DB layer, the backend abstraction, and config — everything needed to persist and retrieve state before any API or SDK work begins.

---

closes: #65777

What's in this PR

1. ORM Models

Two new models: TaskStateModel and AssetStateModel in airflow/models/.

task_state PK is (dag_run_id, task_id, map_index, key):

• dag_run_id is an integer FK to dag_run.id with CASCADE DELETE — task state is automatically wiped when a run is deleted, no orphan cleanup job needed.
• dag_id and run_id are denormalized strings alongside dag_run_id. The integer FK is for the JOIN and CASCADE; the strings are for fast Execution API lookups without a join and for surfacing updated_by_run in the Core API response.
• map_index (default -1) gives each mapped task instance its own namespace. Without it, all instances of the same mapped task share one row per key — instance 0 writes remote_job_id=app_001, instance 1 overwrites it with app_002, instance 0 retries and reconnects to the wrong job.

asset_state PK is (asset_id, key):

• asset_id FK to asset.id with CASCADE DELETE.
• No dag_run_id — asset state is intentionally run-independent. A watermark written in run 1 is readable by run 2. This is what distinguishes it from task state.

Two tables instead of one because the GC lifecycles, index shapes, and FK targets are entirely different for the two scopes. A single discriminated table would mean permanently null columns on every row and would conflate two very different cleanup models.

What the schema implies:

• Retries within the same run share state — same (dag_run_id, task_id, map_index) will have same rows. Try 3 reads remote_job_id written by previous try and reconnects. This is the core AIP-103 use case.
• Each new DAG run starts with a clean task state namespace — no cross-run leakage via task state.
• Null values are not representable — absence is row non-existence. "Not set" and "set to null" are not ambiguous.

2. Alembic Migration

Migration 0112_3_3_0_add_task_state_and_asset_state creates both tables.

3. BaseStateBackend in shared/state

TaskScope, AssetScope, StateScope, and BaseStateBackend live in a new shared distribution (apache-airflow-shared-state) rather than airflow-core. This is the only placement that avoids circular dependencies across all three consumers: provider packages (e.g. S3/GCS custom backends) cannot depend on airflow-core; task-sdk needs the scope dataclasses for context accessors in PR 3; airflow-core needs them for MetastoreStateBackend.

The interface takes a single scope parameter — the backend routes internally on isinstance(scope, TaskScope). All eight methods (four sync, four async) are abstract — implementation must provide real async code. Async variants are required for AIP-98 compatibility so async tasks can call await task_state.aset(...) without blocking the event loop.

The [state] config section (backend key) and resolve_state_backend() follow the resolve_xcom_backend() pattern — validate the configured class is a BaseStateBackend subclass before returning it.

4. MetastoreStateBackend

Default implementation in airflow/state/metastore.py. Routes on scope type, resolves dag_run_id from (dag_id, run_id) before writing task state, and follows the standard @provide_session pattern throughout airflow-core.

---

Next steps

• PR 2: Execution API + Core API endpoints
• PR 3: Task SDK — comms, supervisor, context accessors
• PR 4: GC and retention config

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[amoghrajesh]

@jscheffl@jroachgolf84@ashb do you wanna take another look? Planning to merge it once CI is green otherwise

[amoghrajesh]

Rerunning with full tests

[amoghrajesh]

The failure: https://github.com/apache/airflow/actions/runs/25119239062/job/73621449539?pr=65759 was intermittent, retriggered that job

[amoghrajesh]

Thanks for reviews folks, merging this one in, I handled all comments and resolved it, but if you think you need some changes to be done, let me know, I can take it as a follow up