AIP-103: Adding ability for per task state key retention from operators

[amoghrajesh]

closes: #66461

Why?

The global [state_store] default_retention_days config applies one retention window to every task state key. Some keys have meaningfully different lifetimes — a submitted job ID is useful for the life of a run, while a short-lived lock key might need only hours. This adds the ability to express per-key retention without changing the global default for everything.

Current behaviour

All task state keys written via PUT /state/ti/{id}/{key} receive expires_at = now + default_retention_days, regardless of the individual key's intended lifetime (as established in #66463).

Proposed change

Adds an optional retention: timedelta | None parameter to TaskStateAccessor.set(), along with a public NEVER_EXPIRE constant.

• timedelta(...) — expire this key after the given duration, overriding the global default.
• NEVER_EXPIRE — key never expires, regardless of global config. Stored as NULL in the DB — the GC pass skips NULL rows.
• None / omitted — use [state_store] default_retention_days as before.

The worker computes an absolute expires_at UTC timestamp before sending to the server — the server stores it directly with no date arithmetic.

Changes span the full stack: TaskStateAccessor.set(), SetTaskState comms message, TaskStatePutBody, BaseStateBackend.set() / aset(), MetastoreStateBackend._set_task_state() / _aset_task_state(), and the Execution API route.

User implications / backcompat

None. retention defaults to None. BaseStateBackend.set() and aset() gain a new expires_at: datetime | None = None keyword argument replacing retention_days: int | None = None. expires_at=None means never expire (stored as NULL).

Usage:

fromdatetimeimporttimedeltafromairflow.sdkimportNEVER_EXPIREtask_state.set("job_id", "app_001", retention=timedelta(hours=6))
task_state.set("permanent_key", "value", retention=NEVER_EXPIRE)

Testing

fromdatetimeimporttimedelta, datetimefromairflow.sdkimporttask, DAG, NEVER_EXPIREwithDAG(dag_id="my_dag_for_task_state_retention", schedule=None, start_date=datetime(2022, 3, 4)) asdag:
@taskdeft1(**context):
ts=context["task_state"]
ts.set("short_lived_key", "short_lived_value", retention=timedelta(days=1))
ts.set("shortest_lived_key", "short_lived_value", retention=timedelta(hours=6))
ts.set("long_lived_key", "long_lived_value", retention=timedelta(days=25))
ts.set("never_expire_key", "permanent_value", retention=NEVER_EXPIRE)
ts.set("default_lifetime_key", "default_lifetime_value")
t1()

1. short_lived_key expires 1 day from execution
2. shortest_lived_key expires 6 hours from execution
3. long_lived_key expires 25 days from execution
4. never_expire_key has expires_at = NULL
5. default_lifetime_key uses [state_store] default_retention_days

---

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

• Yes (please specify the tool below)

---

• 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.

[ashb]

I think we should change the Exec API interface -- rather than sending retention_days to the server, lets make the Exec API endpoint(s?) take an expires_at as a DateTime and do any calculation on the client side. This is more expressive, and also lets us much more easily have expiry of less than 1 day. (For instance, for dags that run every 5 or 30 min, 1 day can still be a "long time").

I'm wondering if instead of retention_days being a) days, b) an integer, it should be a timeinterval? WDYT?

It could be retention: int | timedelta | None = None -- int = days, or a timedelta. This might be "too magic", but I think we already have a similar convention for this int=days somewhere in Airflow?

[amoghrajesh]

Thanks for review folks, merging this one.