Skip to content

Add session param to BaseStateBackend interface to fix custom backends - #66708

Merged
amoghrajesh merged 7 commits into
apache:mainfrom
astronomer:worktree-fix-state-backend-session
May 14, 2026
Merged

Add session param to BaseStateBackend interface to fix custom backends#66708
amoghrajesh merged 7 commits into
apache:mainfrom
astronomer:worktree-fix-state-backend-session

Conversation

@amoghrajesh

Copy link
Copy Markdown
Contributor

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Every call site in the execution API for task state and asset state passes session=session to get_state_backend().get/set/delete/clear(...). This is required so the TI scope look-up and the state write share the same sqla transaction, avoiding a race condition where the TI could be deleted between the two operations (flagged in #66073). However, BaseStateBackend had no session parameter — the calls only worked at runtime because @provide_session on MetastoreStateBackend intercepts the kwarg. Any custom backend would crash with TypeError: unexpected keyword argument 'session'.

Current behaviour

All twelve state backend call sites carry # type: ignore[call-arg] comments to silence mypy. A custom backend implementing BaseStateBackend would raise TypeError at runtime the moment any route called get/set/delete/clear on it.

Proposed change

Adds session: Session | None = None to get, set, delete, and clear on BaseStateBackend. The Session type is imported under TYPE_CHECKING only, so there is no runtime SQLAlchemy dependency on the shared airflow_shared.state library. All twelve # type: ignore[call-arg] comments are removed.

Alternatives considered:

  • Remove session=session at call sites — breaks transaction atomicity between scope look-up and write (Kaxil's concern in AIP-103: Add Execution API endpoints for task and asset states #66073).
  • **kwargs — completely opaque; mypy cannot help and custom backend authors have no idea what can be passed.
  • isinstance check at every call site — ugly, defeats the abstraction, requires importing MetastoreStateBackend into route layer.

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

Lee-W
Lee-W previously approved these changes May 11, 2026
Comment threadshared/state/src/airflow_shared/state/__init__.py Outdated
@Lee-W
Lee-W dismissed their stale reviewMay 11, 2026 10:48

overall idea LGTM, but one change needed

Co-authored-by: Wei Lee <weilee.rx@gmail.com>
@amoghrajesh
amoghrajesh requested a review from Lee-WMay 11, 2026 11:05
Comment threadshared/state/src/airflow_shared/state/__init__.py Outdated
@amoghrajesh
amoghrajesh requested a review from ashbMay 12, 2026 08:29
Comment threadshared/state/src/airflow_shared/state/__init__.py Outdated
Comment threadairflow-core/src/airflow/state/metastore.py Outdated
@amoghrajesh

Copy link
Copy Markdown
ContributorAuthor

Failing checks are unrelated, merging this one in.

@amoghrajesh
amoghrajesh merged commit 7e22395 into apache:mainMay 14, 2026
119 of 121 checks passed
@amoghrajesh
amoghrajesh deleted the worktree-fix-state-backend-session branch May 14, 2026 05:03
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

4 participants

@amoghrajesh@ashb@uranusjr@Lee-W