Fix retry policy overrides not persisted to task instance history

[vatsrahul1001]

What's wrong

AIP-105 lets a retry policy choose a custom retry delay and reason each time a task fails. Those are meant to be recorded per-try in task_instance_history as a durable audit trail — "why did the policy wait N seconds on try 3, and for what reason?"

Currently TIH table is always empty: retry_delay_override and retry_reason come out NULL on every task_instance_history row, even when a policy clearly set them.

Root cause

In the Execution API retry handler, the override + reason are written only into the live-row UPDATE, and the task instance is archived to history (prepare_db_for_next_try() → TaskInstanceHistory.record_ti()) before those values are set on the object being archived. record_ti() snapshots columns off the in-memory task-instance object, so it captures the still-unset (NULL) values.

Fix

Set retry_delay_override / retry_reason on the task instance before it's archived, so record_ti() picks them up. Live-row behaviour and retry timing are unchanged.

Testing

Adds a regression test that a policy-driven retry persists the override + reason into task_instance_history (fails on main, passes with the fix). The existing tests only asserted the live row, which is why this shipped.

Before

After

DAG used for testing

from __future__ import annotations

from datetime import timedelta

from airflow.sdk import DAG, ExceptionRetryPolicy, RetryAction, RetryRule, task

QA_POLICY = ExceptionRetryPolicy(
    rules=[
        RetryRule(
            exception=PermissionError,
            action=RetryAction.FAIL,
            reason="Auth failure, not retryable",
        ),
        RetryRule(
            exception=ConnectionError,
            action=RetryAction.RETRY,
            retry_delay=timedelta(seconds=30),
            reason="Transient, backing off 30s",
        ),
        RetryRule(
            exception=TimeoutError,
            action=RetryAction.RETRY,
            retry_delay=timedelta(minutes=2),
            reason="Parked long so live row keeps override at downgrade snapshot",
        ),
    ],
    default=RetryAction.DEFAULT,
)

with DAG(
    dag_id="aip105_migration_baseline",
    schedule=None,
    catchup=False,
    tags=["qa", "aip105", "migration", "retry-policy"],
):

    @task
    def succeeds():
        return "ok"

    # F-01: PermissionError -> FAIL immediately, no override columns written.
    @task(retries=3, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def fail_fast_on_auth():
        raise PermissionError("403 Forbidden")

    # F-02 / history: ConnectionError -> RETRY 30s. Cycles through retries; each
    # try archived to task_instance_history with override=30, reason set.
    @task(retries=3, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def retry_short_cycle():
        raise ConnectionError("connection reset")

    # Live-row populated data: TimeoutError -> RETRY 2h. After the first failure
    # it sits UP_FOR_RETRY; its LIVE task_instance row holds override=7200 and
    # reason until the next attempt (2h out) — snapshot + downgrade before then.
    @task(retries=5, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def retry_parked_long():
        raise TimeoutError("parked for downgrade snapshot")

    # F-03: unmatched exception -> DEFAULT -> standard retry, override stays NULL.
    @task(retries=2, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def default_path():
        raise ValueError("unmatched error")

    succeeds() >> [fail_fast_on_auth(), retry_short_cycle(), retry_parked_long(), default_path()]

---

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

• Yes — Claude Code (Opus 4.8)

Generated-by: Claude Code (Opus 4.8) following the guidelines

[kaxil]

Lgtm and can be merged. Although unsure if this in TIH would be useful at all for anything. Maybe auditing 🤷‍♂️

[vatsrahul1001]

Lgtm and can be merged. Although unsure if this in TIH would be useful at all for anything. Maybe auditing 🤷‍♂️

Yeah auditing is what was mentioned in docs

[github-actions]

Hi maintainer, this PR was merged without a milestone set.
We've automatically set the milestone to Airflow 3.3.1 based on: backport label targeting v3-3-test
If this milestone is not correct, please update it to the appropriate milestone.

This comment was generated by Milestone Tag Assistant.

[github-actions]

Backport successfully created: v3-3-test

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
✅	v3-3-test