Add session param to BaseStateBackend interface to fix custom backends (#66708)

Author: amoghrajesh

airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state.py

@@ -87,7 +87,7 @@ def get_asset_state_by_name(
 ) -> AssetStateResponse:
     """Get an asset state value by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
-    value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session)
     if value is None:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
@@ -105,7 +105,7 @@ def set_asset_state_by_name(
 ) -> None:
     """Set an asset state value by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
-    get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value, session=session)
 
 
 @router.delete("/by-name/value", status_code=status.HTTP_204_NO_CONTENT)
@@ -116,7 +116,7 @@ def delete_asset_state_by_name(
 ) -> None:
     """Delete a single asset state key by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
-    get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)
 
 
 @router.delete("/by-name/clear", status_code=status.HTTP_204_NO_CONTENT)
@@ -126,7 +126,7 @@ def clear_asset_state_by_name(
 ) -> None:
     """Delete all state keys for an asset by asset name."""
     asset_id = _resolve_asset_id_by_name(name, session)
-    get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)
 
 
 @router.get("/by-uri/value")
@@ -137,7 +137,7 @@ def get_asset_state_by_uri(
 ) -> AssetStateResponse:
     """Get an asset state value by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
-    value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    value = get_state_backend().get(AssetScope(asset_id=asset_id), key, session=session)
     if value is None:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
@@ -155,7 +155,7 @@ def set_asset_state_by_uri(
 ) -> None:
     """Set an asset state value by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
-    get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().set(AssetScope(asset_id=asset_id), key, body.value, session=session)
 
 
 @router.delete("/by-uri/value", status_code=status.HTTP_204_NO_CONTENT)
@@ -166,7 +166,7 @@ def delete_asset_state_by_uri(
 ) -> None:
     """Delete a single asset state key by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
-    get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().delete(AssetScope(asset_id=asset_id), key, session=session)
 
 
 @router.delete("/by-uri/clear", status_code=status.HTTP_204_NO_CONTENT)
@@ -176,4 +176,4 @@ def clear_asset_state_by_uri(
 ) -> None:
     """Delete all state keys for an asset by asset URI."""
     asset_id = _resolve_asset_id_by_uri(uri, session)
-    get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().clear(AssetScope(asset_id=asset_id), session=session)

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state.py

@@ -65,7 +65,7 @@ def get_task_state(
 ) -> TaskStateResponse:
     """Get value for a task state."""
     scope = _get_task_scope_for_ti(task_instance_id, session)
-    value = get_state_backend().get(scope, key, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    value = get_state_backend().get(scope, key, session=session)
     if value is None:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
@@ -86,7 +86,7 @@ def set_task_state(
 ) -> None:
     """Set a task state key, creating or updating the row."""
     scope = _get_task_scope_for_ti(task_instance_id, session)
-    get_state_backend().set(scope, key, body.value, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().set(scope, key, body.value, session=session)
 
 
 @router.delete("/{task_instance_id}/{key}", status_code=status.HTTP_204_NO_CONTENT)
@@ -97,7 +97,7 @@ def delete_task_state(
 ) -> None:
     """Delete a single task state key."""
     scope = _get_task_scope_for_ti(task_instance_id, session)
-    get_state_backend().delete(scope, key, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().delete(scope, key, session=session)
 
 
 @router.delete("/{task_instance_id}", status_code=status.HTTP_204_NO_CONTENT)
@@ -125,4 +125,4 @@ def clear_task_state(
     accepted without error.
     """
     scope = _get_task_scope_for_ti(task_instance_id, session)
-    get_state_backend().clear(scope, all_map_indices=all_map_indices, session=session)  # type: ignore[call-arg]  # @provide_session adds session kwarg at runtime; BaseStateBackend signature omits it so mypy can't see it
+    get_state_backend().clear(scope, all_map_indices=all_map_indices, session=session)

airflow-core/src/airflow/state/metastore.py

@@ -17,6 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING
 
 from sqlalchemy import delete, select
@@ -38,6 +40,16 @@
     from sqlalchemy.orm import Session
 
 
+@asynccontextmanager
+async def _async_session(session: AsyncSession | None) -> AsyncGenerator[AsyncSession, None]:
+    """Use provided async session or create a new one."""
+    if session is not None:
+        yield session
+    else:
+        async with create_session_async() as s:
+            yield s
+
+
 def _build_upsert_stmt(
     dialect: str | None,
     model: type,
@@ -69,7 +81,9 @@ class MetastoreStateBackend(BaseStateBackend):
     """Default state backend for tasks and assets. Stores task and asset state in the Airflow metadata database."""
 
     @provide_session
-    def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> str | None:
+    def get(self, scope: StateScope, key: str, *, session: Session | None = NEW_SESSION) -> str | None:
+        if TYPE_CHECKING:
+            assert session is not None
         match scope:
             case TaskScope():
                 return self._get_task_state(scope, key, session=session)
@@ -79,7 +93,9 @@ def get(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) ->
                 assert_never(scope)
 
     @provide_session
-    def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None:
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session | None = NEW_SESSION) -> None:
+        if TYPE_CHECKING:
+            assert session is not None
         match scope:
             case TaskScope():
                 self._set_task_state(scope, key, value, session=session)
@@ -89,7 +105,9 @@ def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW
                 assert_never(scope)
 
     @provide_session
-    def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None:
+    def delete(self, scope: StateScope, key: str, *, session: Session | None = NEW_SESSION) -> None:
+        if TYPE_CHECKING:
+            assert session is not None
         match scope:
             case TaskScope():
                 self._delete_task_state(scope, key, session=session)
@@ -104,8 +122,10 @@ def clear(
         scope: StateScope,
         *,
         all_map_indices: bool = False,
-        session: Session = NEW_SESSION,
+        session: Session | None = NEW_SESSION,
     ) -> None:
+        if TYPE_CHECKING:
+            assert session is not None
         match scope:
             case TaskScope():
                 self._clear_task_state(scope, all_map_indices=all_map_indices, session=session)
@@ -114,43 +134,47 @@ def clear(
             case _:
                 assert_never(scope)
 
-    async def aget(self, scope: StateScope, key: str) -> str | None:
-        async with create_session_async() as session:
+    async def aget(self, scope: StateScope, key: str, *, session: AsyncSession | None = None) -> str | None:
+        async with _async_session(session) as s:
             match scope:
                 case TaskScope():
-                    return await self._aget_task_state(scope, key, session=session)
+                    return await self._aget_task_state(scope, key, session=s)
                 case AssetScope():
-                    return await self._aget_asset_state(scope, key, session=session)
+                    return await self._aget_asset_state(scope, key, session=s)
                 case _:
                     assert_never(scope)
 
-    async def aset(self, scope: StateScope, key: str, value: str) -> None:
-        async with create_session_async() as session:
+    async def aset(
+        self, scope: StateScope, key: str, value: str, *, session: AsyncSession | None = None
+    ) -> None:
+        async with _async_session(session) as s:
             match scope:
                 case TaskScope():
-                    await self._aset_task_state(scope, key, value, session=session)
+                    await self._aset_task_state(scope, key, value, session=s)
                 case AssetScope():
-                    await self._aset_asset_state(scope, key, value, session=session)
+                    await self._aset_asset_state(scope, key, value, session=s)
                 case _:
                     assert_never(scope)
 
-    async def adelete(self, scope: StateScope, key: str) -> None:
-        async with create_session_async() as session:
+    async def adelete(self, scope: StateScope, key: str, *, session: AsyncSession | None = None) -> None:
+        async with _async_session(session) as s:
             match scope:
                 case TaskScope():
-                    await self._adelete_task_state(scope, key, session=session)
+                    await self._adelete_task_state(scope, key, session=s)
                 case AssetScope():
-                    await self._adelete_asset_state(scope, key, session=session)
+                    await self._adelete_asset_state(scope, key, session=s)
                 case _:
                     assert_never(scope)
 
-    async def aclear(self, scope: StateScope, *, all_map_indices: bool = False) -> None:
-        async with create_session_async() as session:
+    async def aclear(
+        self, scope: StateScope, *, all_map_indices: bool = False, session: AsyncSession | None = None
+    ) -> None:
+        async with _async_session(session) as s:
             match scope:
                 case TaskScope():
-                    await self._aclear_task_state(scope, all_map_indices=all_map_indices, session=session)
+                    await self._aclear_task_state(scope, all_map_indices=all_map_indices, session=s)
                 case AssetScope():
-                    await self._aclear_asset_state(scope, session=session)
+                    await self._aclear_asset_state(scope, session=s)
                 case _:
                     assert_never(scope)
 

airflow-core/tests/unit/state/test_metastore.py

@@ -28,7 +28,7 @@
 from airflow.models.task_state import TaskStateModel
 from airflow.state import AssetScope, TaskScope, resolve_state_backend
 from airflow.state.metastore import MetastoreStateBackend
-from airflow.utils.session import create_session
+from airflow.utils.session import create_session, create_session_async
 
 from tests_common.test_utils.config import conf_vars
 from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs
@@ -379,6 +379,16 @@ async def test_aset_task_raises_for_missing_dag_run(self, backend: MetastoreStat
         with pytest.raises(ValueError, match="No DagRun found"):
             await backend.aset(scope, "job_id", "app_async")
 
+    async def test_aset_and_aget_with_provided_session(
+        self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+    ):
+        """async methods use a provided AsyncSession when one is given."""
+        scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID)
+        async with create_session_async() as session:
+            await backend.aset(scope, "job_id", "app_with_session", session=session)
+            result = await backend.aget(scope, "job_id", session=session)
+        assert result == "app_with_session"
+
 
 class TestResolveStateBackend:
     @conf_vars({("state_store", "backend"): "airflow.state.metastore.MetastoreStateBackend"})

shared/state/src/airflow_shared/state/__init__.py

@@ -18,6 +18,11 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+    from sqlalchemy.orm import Session
 
 
 @dataclass(frozen=True)
@@ -62,34 +67,42 @@ class BaseStateBackend(ABC):
                 ...  # asset-specific storage
 
     Custom backends are configured via ``[state_store] backend`` in ``airflow.cfg``.
+
+    **The ``session`` parameter on ``get``, ``set``, ``delete``, and ``clear``:**
+
+    The default ``MetastoreStateBackend`` passes a SQLAlchemy ``Session`` through
+    these methods. Custom backends that do not use SQLAlchemy should accept ``session`` as a
+    keyword argument and ignore it.
     """
 
     @abstractmethod
-    def get(self, scope: StateScope, key: str) -> str | None:
+    def get(self, scope: StateScope, key: str, *, session: Session | None = None) -> str | None:
         """
         Return the stored value, or None if the key does not exist.
 
         Must handle both ``TaskScope`` and ``AssetScope``.
         """
 
     @abstractmethod
-    def set(self, scope: StateScope, key: str, value: str) -> None:
+    def set(self, scope: StateScope, key: str, value: str, *, session: Session | None = None) -> None:
         """
         Write or overwrite the value for the given key.
 
         Must handle both ``TaskScope`` and ``AssetScope``.
         """
 
     @abstractmethod
-    def delete(self, scope: StateScope, key: str) -> None:
+    def delete(self, scope: StateScope, key: str, *, session: Session | None = None) -> None:
         """
         Delete a single key. No-op if the key does not exist.
 
         Must handle both ``TaskScope`` and ``AssetScope``.
         """
 
     @abstractmethod
-    def clear(self, scope: StateScope, *, all_map_indices: bool = False) -> None:
+    def clear(
+        self, scope: StateScope, *, all_map_indices: bool = False, session: Session | None = None
+    ) -> None:
         """
         Delete all keys under the given scope.
 
@@ -102,23 +115,45 @@ def clear(self, scope: StateScope, *, all_map_indices: bool = False) -> None:
         """
 
     @abstractmethod
-    async def aget(self, scope: StateScope, key: str) -> str | None:
-        """Async variant of get. Must handle both ``TaskScope`` and ``AssetScope``."""
+    async def aget(self, scope: StateScope, key: str, *, session: AsyncSession | None = None) -> str | None:
+        """
+        Async variant of get. Must handle both ``TaskScope`` and ``AssetScope``.
+
+        ``session`` is optional. If provided, implementations should use it directly.
+        If ``None``, implementations manage their own async session internally.
+        """
 
     @abstractmethod
-    async def aset(self, scope: StateScope, key: str, value: str) -> None:
-        """Async variant of set. Must handle both ``TaskScope`` and ``AssetScope``."""
+    async def aset(
+        self, scope: StateScope, key: str, value: str, *, session: AsyncSession | None = None
+    ) -> None:
+        """
+        Async variant of set. Must handle both ``TaskScope`` and ``AssetScope``.
+
+        ``session`` is optional. If provided, implementations should use it directly.
+        If ``None``, implementations manage their own async session internally.
+        """
 
     @abstractmethod
-    async def adelete(self, scope: StateScope, key: str) -> None:
-        """Async variant of delete. Must handle both ``TaskScope`` and ``AssetScope``."""
+    async def adelete(self, scope: StateScope, key: str, *, session: AsyncSession | None = None) -> None:
+        """
+        Async variant of delete. Must handle both ``TaskScope`` and ``AssetScope``.
+
+        ``session`` is optional. If provided, implementations should use it directly.
+        If ``None``, implementations manage their own async session internally.
+        """
 
     @abstractmethod
-    async def aclear(self, scope: StateScope, *, all_map_indices: bool = False) -> None:
+    async def aclear(
+        self, scope: StateScope, *, all_map_indices: bool = False, session: AsyncSession | None = None
+    ) -> None:
         """
         Async variant of clear. Must handle both ``TaskScope`` and ``AssetScope``.
 
         For ``TaskScope``: by default, only keys for the exact ``map_index`` on the
         scope are cleared. Pass ``all_map_indices=True`` to wipe state across every
         mapped instance of the task. For ``AssetScope`` the flag has no effect.
+
+        ``session`` is optional. If provided, implementations should use it directly.
+        If ``None``, implementations manage their own async session internally.
         """