amoghrajesh opened a new pull request, #65759:
URL: https://github.com/apache/airflow/pull/65759
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [x] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
# 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.
---
## 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.
- `updated_by_try` is nullable — it has no meaning for asset state and is
absent there.
`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: 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)`. Async variants (`aget`, `aset`,
`adelete`, `aclear`) are provided as defaults that delegate to the sync
methods, so sync backends get async support for free. These are required for
AIP-98 compatibility — async tasks must be able to 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` / `NEW_SESSION` pattern used
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](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
