kaxil commented on code in PR #65759: URL: https://github.com/apache/airflow/pull/65759#discussion_r3153690860
########## airflow-core/src/airflow/state/metastore.py: ########## @@ -0,0 +1,334 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sqlalchemy import delete, select + +from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope +from airflow._shared.timezones import timezone +from airflow.models.asset_state import AssetStateModel +from airflow.models.dagrun import DagRun +from airflow.models.task_state import TaskStateModel +from airflow.utils.session import NEW_SESSION, create_session_async, provide_session +from airflow.utils.sqlalchemy import get_dialect_name + +if TYPE_CHECKING: + from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert + from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert + from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert + from sqlalchemy.ext.asyncio import AsyncSession + from sqlalchemy.orm import Session + + +def _build_upsert_stmt( + dialect: str | None, + model: type, + conflict_cols: list[str], + values: dict, + update_fields: dict, +): + """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement.""" + stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert + if dialect == "postgresql": + from sqlalchemy.dialects.postgresql import insert as pg_insert + + stmt = pg_insert(model).values(**values) + stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields) + elif dialect == "mysql": + from sqlalchemy.dialects.mysql import insert as mysql_insert + + stmt = mysql_insert(model).values(**values) + stmt = stmt.on_duplicate_key_update(**update_fields) + else: + from sqlalchemy.dialects.sqlite import insert as sqlite_insert + + stmt = sqlite_insert(model).values(**values) + stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields) + return stmt + + +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: + match scope: + case TaskScope(): + return self._get_task_state(scope, key, session=session) + case AssetScope(): + return self._get_asset_state(scope, key, session=session) + + @provide_session + def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None: + match scope: + case TaskScope(): + self._set_task_state(scope, key, value, session=session) + case AssetScope(): + self._set_asset_state(scope, key, value, session=session) + + @provide_session + def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None: + match scope: + case TaskScope(): + self._delete_task_state(scope, key, session=session) + case AssetScope(): + self._delete_asset_state(scope, key, session=session) + + @provide_session + def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None: + match scope: + case TaskScope(): + self._clear_task_state(scope, session=session) + case AssetScope(): + self._clear_asset_state(scope, session=session) + + async def aget(self, scope: StateScope, key: str) -> str | None: + async with create_session_async() as session: + match scope: + case TaskScope(): + return await self._aget_task_state(scope, key, session=session) + case AssetScope(): + return await self._aget_asset_state(scope, key, session=session) + + async def aset(self, scope: StateScope, key: str, value: str) -> None: + async with create_session_async() as session: + match scope: + case TaskScope(): + await self._aset_task_state(scope, key, value, session=session) + case AssetScope(): + await self._aset_asset_state(scope, key, value, session=session) + + async def adelete(self, scope: StateScope, key: str) -> None: + async with create_session_async() as session: + match scope: + case TaskScope(): + await self._adelete_task_state(scope, key, session=session) + case AssetScope(): + await self._adelete_asset_state(scope, key, session=session) + + async def aclear(self, scope: StateScope) -> None: + async with create_session_async() as session: + match scope: + case TaskScope(): + await self._aclear_task_state(scope, session=session) + case AssetScope(): + await self._aclear_asset_state(scope, session=session) + + def _get_task_state(self, scope: TaskScope, key: str, *, session: Session) -> str | None: + row = session.scalar( + select(TaskStateModel).where( + TaskStateModel.dag_id == scope.dag_id, + TaskStateModel.run_id == scope.run_id, + TaskStateModel.task_id == scope.task_id, + TaskStateModel.map_index == scope.map_index, + TaskStateModel.key == key, + ) + ) + return row.value if row is not None else None + + def _set_task_state(self, scope: TaskScope, key: str, value: str, *, session: Session) -> None: + dag_run_id = session.scalar( + select(DagRun.id).where( + DagRun.dag_id == scope.dag_id, + DagRun.run_id == scope.run_id, + ) + ) + if dag_run_id is None: + raise ValueError(f"No DagRun found for dag_id={scope.dag_id!r} run_id={scope.run_id!r}") + now = timezone.utcnow() + values = dict( + dag_run_id=dag_run_id, + dag_id=scope.dag_id, + run_id=scope.run_id, + task_id=scope.task_id, + map_index=scope.map_index, + key=key, + value=value, + updated_at=now, + ) + stmt = _build_upsert_stmt( + get_dialect_name(session), + TaskStateModel, + ["dag_run_id", "task_id", "map_index", "key"], + values, + dict(value=value, updated_at=now), + ) + session.execute(stmt) + + def _delete_task_state(self, scope: TaskScope, key: str, *, session: Session) -> None: + session.execute( + delete(TaskStateModel).where( + TaskStateModel.dag_id == scope.dag_id, + TaskStateModel.run_id == scope.run_id, + TaskStateModel.task_id == scope.task_id, + TaskStateModel.map_index == scope.map_index, + TaskStateModel.key == key, + ) + ) + + def _clear_task_state(self, scope: TaskScope, *, session: Session) -> None: Review Comment: `_clear_task_state` filters `WHERE map_index = scope.map_index`, and `TaskScope.map_index` defaults to `-1`. So `backend.clear(TaskScope(dag_id, run_id, task_id))` only deletes non-mapped rows -- for a mapped task with state at map_index 0/1/2, this is a silent no-op. The new `test_clear_scoped_to_map_index` actually demonstrates the surprise (clear at index 0 leaves index 1 alone), but there's no path to clear all map indices for a task. Either drop the default on `map_index`, or special-case `-1` inside `clear`/`aclear` to mean "all map indices for this task", or document the limit on both `TaskScope` and `clear()`. ########## airflow-core/tests/unit/state/test_metastore.py: ########## @@ -0,0 +1,363 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from sqlalchemy import select + +from airflow._shared.timezones import timezone +from airflow.models.asset import AssetModel +from airflow.models.dagrun import DagRun, DagRunType +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 tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + +pytestmark = pytest.mark.db_test + +DAG_ID = "test_dag" +TASK_ID = "test_task" +RUN_ID = "scheduled__2026-04-24" + + [email protected](autouse=True) +def clean_tables(): + clear_db_dags() + clear_db_runs() + clear_db_assets() + yield + clear_db_dags() + clear_db_runs() + clear_db_assets() + + [email protected] +def backend() -> MetastoreStateBackend: + return MetastoreStateBackend() + + [email protected] +def dag_run(session: Session) -> DagRun: + run = DagRun( + dag_id=DAG_ID, + run_id=RUN_ID, + run_type=DagRunType.SCHEDULED, + logical_date=timezone.datetime(2026, 4, 24), + run_after=timezone.datetime(2026, 4, 24), + ) + session.add(run) + session.flush() + return run + + [email protected] +def dag_run_committed() -> DagRun: + """DagRun committed to DB so async sessions (which open their own connections) can see it.""" + with create_session() as session: + run = DagRun( + dag_id=DAG_ID, + run_id=RUN_ID, + run_type=DagRunType.SCHEDULED, + logical_date=timezone.datetime(2026, 4, 24), + run_after=timezone.datetime(2026, 4, 24), + ) + session.add(run) + session.flush() + session.expunge(run) + return run + + [email protected] +def asset(session: Session) -> AssetModel: + a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test") + session.add(a) + session.flush() + return a + + [email protected] +def asset_committed() -> AssetModel: + """AssetModel committed to DB so async sessions (which open their own connections) can see it.""" + with create_session() as session: + a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test") + session.add(a) + session.flush() + session.expunge(a) + return a + + +class TestMetastoreStateBackendTaskScope: + def test_get_returns_none_for_missing_key( + self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun + ): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + assert backend.get(scope, "missing", session=session) is None + + def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + backend.set(scope, "job_id", "app_1234", session=session) + session.flush() + assert backend.get(scope, "job_id", session=session) == "app_1234" + + def test_set_twice_overrides_existing_value( + self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun + ): + """Calling set twice on the same key updates the value in place.""" + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + backend.set(scope, "job_id", "app_try1", session=session) + session.flush() + backend.set(scope, "job_id", "app_try2", session=session) + session.flush() + assert backend.get(scope, "job_id", session=session) == "app_try2" + + def test_set_stores_dag_run_id_fk( + self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun + ): + """set resolves dag_run_id from (dag_id, run_id) and persists it as the FK.""" + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + backend.set(scope, "job_id", "app_1234", session=session) + session.flush() + + row = session.scalar( + select(TaskStateModel).where( + TaskStateModel.dag_id == DAG_ID, + TaskStateModel.task_id == TASK_ID, + TaskStateModel.key == "job_id", + ) + ) + assert row is not None + assert row.dag_run_id == dag_run.id + + def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + backend.set(scope, "job_id", "app_1234", session=session) + backend.set(scope, "checkpoint", "step_3", session=session) + session.flush() + + backend.delete(scope, "job_id", session=session) + session.flush() + + assert backend.get(scope, "job_id", session=session) is None + assert backend.get(scope, "checkpoint", session=session) == "step_3" + + def test_delete_is_noop_for_missing_key( + self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun + ): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + backend.delete(scope, "nonexistent", session=session) + + def test_clear_removes_all_keys(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun): + """clear removes every key under the given task scope.""" + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + for key in ("job_id", "checkpoint", "watermark"): + backend.set(scope, key, f"val_{key}", session=session) + session.flush() + + backend.clear(scope, session=session) + session.flush() + + for key in ("job_id", "checkpoint", "watermark"): + assert backend.get(scope, key, session=session) is None + + def test_map_index_isolation(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun): + """Mapped task instances with different map_index values have isolated namespaces.""" + scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0) + scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1) + + backend.set(scope0, "job_id", "app_0", session=session) + backend.set(scope1, "job_id", "app_1", session=session) + session.flush() + + assert backend.get(scope0, "job_id", session=session) == "app_0" + assert backend.get(scope1, "job_id", session=session) == "app_1" + + def test_set_raises_for_missing_dag_run(self, session: Session, backend: MetastoreStateBackend): + scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID) + with pytest.raises(ValueError, match="No DagRun found"): + backend.set(scope, "job_id", "app_1234", session=session) + + def test_clear_scoped_to_map_index( + self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun + ): + """clear for map_index=0 does not remove state belonging to map_index=1.""" + scope0 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=0) + scope1 = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID, map_index=1) + + backend.set(scope0, "job_id", "app_0", session=session) + backend.set(scope1, "job_id", "app_1", session=session) + session.flush() + + backend.clear(scope0, session=session) + session.flush() + + assert backend.get(scope0, "job_id", session=session) is None + assert backend.get(scope1, "job_id", session=session) == "app_1" + + +class TestMetastoreStateBackendAssetScope: + def test_get_returns_none_for_missing_key( + self, session: Session, backend: MetastoreStateBackend, asset: AssetModel + ): + scope = AssetScope(asset_id=asset.id) + assert backend.get(scope, "missing", session=session) is None + + def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel): + scope = AssetScope(asset_id=asset.id) + backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session) + session.flush() + assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z" + + def test_set_twice_overwrites_existing_value( + self, session: Session, backend: MetastoreStateBackend, asset: AssetModel + ): + scope = AssetScope(asset_id=asset.id) + backend.set(scope, "watermark", "2026-04-01T00:00:00Z", session=session) + session.flush() + backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session) + session.flush() + assert backend.get(scope, "watermark", session=session) == "2026-04-24T00:00:00Z" + + def test_delete_removes_key(self, session: Session, backend: MetastoreStateBackend, asset: AssetModel): + scope = AssetScope(asset_id=asset.id) + backend.set(scope, "watermark", "2026-04-24T00:00:00Z", session=session) + backend.set(scope, "file_count", "42", session=session) + session.flush() + + backend.delete(scope, "watermark", session=session) + session.flush() + + assert backend.get(scope, "watermark", session=session) is None + assert backend.get(scope, "file_count", session=session) == "42" + + def test_delete_is_noop_for_missing_key( + self, session: Session, backend: MetastoreStateBackend, asset: AssetModel + ): + scope = AssetScope(asset_id=asset.id) + backend.delete(scope, "nonexistent", session=session) + + def test_clear_removes_all_keys( + self, session: Session, backend: MetastoreStateBackend, asset: AssetModel + ): + scope = AssetScope(asset_id=asset.id) + for key in ("watermark", "file_count", "last_error"): + backend.set(scope, key, f"val_{key}", session=session) + session.flush() + + backend.clear(scope, session=session) + session.flush() + + for key in ("watermark", "file_count", "last_error"): + assert backend.get(scope, key, session=session) is None + + def test_different_assets_are_isolated( + self, session: Session, backend: MetastoreStateBackend, asset: AssetModel + ): + asset2 = AssetModel(uri="s3://bucket/other", name="other_asset", group="test") + session.add(asset2) + session.flush() + + scope1 = AssetScope(asset_id=asset.id) + scope2 = AssetScope(asset_id=asset2.id) + + backend.set(scope1, "watermark", "asset1_value", session=session) + session.flush() + + assert backend.get(scope2, "watermark", session=session) is None + + [email protected](loop_scope="class") +class TestMetastoreStateBackendAsync: + async def test_aset_and_aget_task_roundtrip( + self, backend: MetastoreStateBackend, dag_run_committed: DagRun + ): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + await backend.aset(scope, "job_id", "app_async") + result = await backend.aget(scope, "job_id") + assert result == "app_async" + + async def test_adelete_task_removes_key(self, backend: MetastoreStateBackend, dag_run_committed: DagRun): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + await backend.aset(scope, "job_id", "app_async") + await backend.adelete(scope, "job_id") + assert await backend.aget(scope, "job_id") is None + + async def test_aclear_task_removes_all_keys( + self, backend: MetastoreStateBackend, dag_run_committed: DagRun + ): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + await backend.aset(scope, "job_id", "app_async") + await backend.aset(scope, "checkpoint", "step_1") + await backend.aclear(scope) + assert await backend.aget(scope, "job_id") is None + assert await backend.aget(scope, "checkpoint") is None + + async def test_aset_and_aget_asset_roundtrip( + self, backend: MetastoreStateBackend, asset_committed: AssetModel + ): + scope = AssetScope(asset_id=asset_committed.id) + await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z") + result = await backend.aget(scope, "watermark") + assert result == "2026-04-27T00:00:00Z" + + async def test_adelete_asset_removes_key( + self, backend: MetastoreStateBackend, asset_committed: AssetModel + ): + scope = AssetScope(asset_id=asset_committed.id) + await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z") + await backend.adelete(scope, "watermark") + assert await backend.aget(scope, "watermark") is None + + async def test_aclear_asset_removes_all_keys( + self, backend: MetastoreStateBackend, asset_committed: AssetModel + ): + scope = AssetScope(asset_id=asset_committed.id) + await backend.aset(scope, "watermark", "2026-04-27T00:00:00Z") + await backend.aset(scope, "file_count", "42") + await backend.aclear(scope) + assert await backend.aget(scope, "watermark") is None + assert await backend.aget(scope, "file_count") is None + + async def test_aset_task_raises_for_missing_dag_run(self, backend: MetastoreStateBackend): + scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run", task_id=TASK_ID) + with pytest.raises(ValueError, match="No DagRun found"): + await backend.aset(scope, "job_id", "app_async") + + +class TestResolveStateBackend: + def test_default_resolves_to_metastore_backend(self): Review Comment: `test_default_resolves_to_metastore_backend` calls `resolve_state_backend()` with no `conf_vars` override, so it depends on global config state -- if a previous test sets `state.backend = ""` and forgets cleanup, this fails non-deterministically. It also passes by accident: it's exercising the `config.yml` default rather than the `if not clazz` branch in `resolve_state_backend`. Add a separate test wrapped in `conf_vars({("state", "backend"): ""})` that asserts the explicit-empty path (which today returns `MetastoreStateBackend`, but the previous review suggested raising instead -- pick one and lock it in), and rename this one to `test_resolve_returns_default_from_config_yml`. ########## airflow-core/tests/unit/state/test_metastore.py: ########## @@ -0,0 +1,363 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from sqlalchemy import select + +from airflow._shared.timezones import timezone +from airflow.models.asset import AssetModel +from airflow.models.dagrun import DagRun, DagRunType +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 tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import clear_db_assets, clear_db_dags, clear_db_runs + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + +pytestmark = pytest.mark.db_test + +DAG_ID = "test_dag" +TASK_ID = "test_task" +RUN_ID = "scheduled__2026-04-24" + + [email protected](autouse=True) +def clean_tables(): + clear_db_dags() + clear_db_runs() + clear_db_assets() + yield + clear_db_dags() + clear_db_runs() + clear_db_assets() + + [email protected] +def backend() -> MetastoreStateBackend: + return MetastoreStateBackend() + + [email protected] +def dag_run(session: Session) -> DagRun: + run = DagRun( + dag_id=DAG_ID, + run_id=RUN_ID, + run_type=DagRunType.SCHEDULED, + logical_date=timezone.datetime(2026, 4, 24), + run_after=timezone.datetime(2026, 4, 24), + ) + session.add(run) + session.flush() + return run + + [email protected] +def dag_run_committed() -> DagRun: + """DagRun committed to DB so async sessions (which open their own connections) can see it.""" + with create_session() as session: + run = DagRun( + dag_id=DAG_ID, + run_id=RUN_ID, + run_type=DagRunType.SCHEDULED, + logical_date=timezone.datetime(2026, 4, 24), + run_after=timezone.datetime(2026, 4, 24), + ) + session.add(run) + session.flush() + session.expunge(run) + return run + + [email protected] +def asset(session: Session) -> AssetModel: + a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test") + session.add(a) + session.flush() + return a + + [email protected] +def asset_committed() -> AssetModel: + """AssetModel committed to DB so async sessions (which open their own connections) can see it.""" + with create_session() as session: + a = AssetModel(uri="s3://bucket/prefix", name="test_asset", group="test") + session.add(a) + session.flush() + session.expunge(a) + return a + + +class TestMetastoreStateBackendTaskScope: + def test_get_returns_none_for_missing_key( + self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun + ): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + assert backend.get(scope, "missing", session=session) is None + + def test_set_and_get_roundtrip(self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun): + scope = TaskScope(dag_id=DAG_ID, run_id=RUN_ID, task_id=TASK_ID) + backend.set(scope, "job_id", "app_1234", session=session) + session.flush() + assert backend.get(scope, "job_id", session=session) == "app_1234" + + def test_set_twice_overrides_existing_value( Review Comment: These set/get/upsert tests only run against the default backend in CI (sqlite). A typo in the `pg_insert.on_conflict_do_update(...)` or `mysql_insert.on_duplicate_key_update(...)` branches in `_build_upsert_stmt` wouldn't be caught -- and given the dialect routing is the load-bearing piece for this PR, that's a real gap. Add a `pytest.mark.backend("postgres", "mysql", "sqlite")` parametrisation on at least `test_set_twice_overrides_existing_value` so each upsert branch gets exercised under its real dialect. ########## airflow-core/src/airflow/config_templates/config.yml: ########## @@ -2997,6 +2997,20 @@ dag_processor: example: ~ default: "warning" +state: Review Comment: Re-flagging the section-name concern from the prior round. `state` is heavily overloaded in Airflow (`TI.state`, `DR.state`, `TaskInstanceState`, `state_color_*`, etc.), so users will read `[state]` and reasonably expect it to be about TI/DR state behaviour, not a key/value blob backend. `[state_store]` parallels `[secrets]` and reads cleanly; `[task_state]` also works. Worth picking now since AIP-103 will likely add per-backend keys (retention, GC) in the next PRs and renaming a config section after release is painful. ########## airflow-core/src/airflow/state/metastore.py: ########## @@ -0,0 +1,334 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sqlalchemy import delete, select + +from airflow._shared.state import AssetScope, BaseStateBackend, StateScope, TaskScope +from airflow._shared.timezones import timezone +from airflow.models.asset_state import AssetStateModel +from airflow.models.dagrun import DagRun +from airflow.models.task_state import TaskStateModel +from airflow.utils.session import NEW_SESSION, create_session_async, provide_session +from airflow.utils.sqlalchemy import get_dialect_name + +if TYPE_CHECKING: + from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert + from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert + from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert + from sqlalchemy.ext.asyncio import AsyncSession + from sqlalchemy.orm import Session + + +def _build_upsert_stmt( + dialect: str | None, + model: type, + conflict_cols: list[str], + values: dict, + update_fields: dict, +): + """Return a dialect-specific INSERT ... ON CONFLICT UPDATE statement.""" + stmt: MySQLInsert | PostgreSQLInsert | SQLiteInsert + if dialect == "postgresql": + from sqlalchemy.dialects.postgresql import insert as pg_insert + + stmt = pg_insert(model).values(**values) + stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields) + elif dialect == "mysql": + from sqlalchemy.dialects.mysql import insert as mysql_insert + + stmt = mysql_insert(model).values(**values) + stmt = stmt.on_duplicate_key_update(**update_fields) + else: + from sqlalchemy.dialects.sqlite import insert as sqlite_insert + + stmt = sqlite_insert(model).values(**values) + stmt = stmt.on_conflict_do_update(index_elements=conflict_cols, set_=update_fields) + return stmt + + +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: + match scope: Review Comment: The `match scope:` blocks here and in the seven sibling methods (lines 80, 88, 96, 104, 112, 120, 128) have no `case _:` arm. `StateScope = TaskScope | AssetScope` today, but adding a future scope subtype makes `get`/`aget` silently return `None` and `set`/`delete`/`clear` silently no-op (writes vanish, no log, no error). Add `case _: assert_never(scope)` (from `typing`) to all eight match blocks so mypy enforces exhaustiveness on every new subtype and runtime fails loudly. This also doubles as the dispatch documentation flagged in the previous review. ########## shared/state/tests/state/test_state.py: ########## @@ -0,0 +1,84 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio + +import pytest + +from airflow_shared.state import BaseStateBackend, StateScope, TaskScope + + +class TestBaseStateBackend: + @pytest.fixture + def backend(self) -> BaseStateBackend: + class ConcreteBackend(BaseStateBackend): + def __init__(self): + self.calls: list[str] = [] + + def get(self, scope: StateScope, key: str) -> str | None: + self.calls.append("get") + return "value" + + def set(self, scope: StateScope, key: str, value: str) -> None: + self.calls.append("set") + + def delete(self, scope: StateScope, key: str) -> None: + self.calls.append("delete") + + def clear(self, scope: StateScope) -> None: + self.calls.append("clear") + + async def aget(self, scope: StateScope, key: str) -> str | None: + self.calls.append("aget") + return "value" + + async def aset(self, scope: StateScope, key: str, value: str) -> None: + self.calls.append("aset") + + async def adelete(self, scope: StateScope, key: str) -> None: + self.calls.append("adelete") + + async def aclear(self, scope: StateScope) -> None: + self.calls.append("aclear") + + return ConcreteBackend() + + def test_aget_is_callable(self, backend): Review Comment: `test_aget_is_callable` and the three siblings (`test_aset_is_callable`, `test_adelete_is_callable`, `test_aclear_is_callable`) test that `asyncio.run` runs an `async def` and that `list.append` works -- none of them exercise anything specific to `BaseStateBackend`. They'll keep passing if you delete every async method from the ABC, or replace them with sync ones. Replace with: (a) a test that instantiating an incomplete `BaseStateBackend` subclass raises `TypeError` ("abstract"), and (b) one that asserts `BaseStateBackend.__abstractmethods__` contains every `a*` method paired with its sync sibling -- that's the actual contract this file should be defending. ########## airflow-core/src/airflow/models/task_state.py: ########## @@ -0,0 +1,62 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from datetime import datetime + +from sqlalchemy import ForeignKeyConstraint, Index, Integer, PrimaryKeyConstraint, String, Text +from sqlalchemy.dialects.mysql import MEDIUMTEXT +from sqlalchemy.orm import Mapped, mapped_column + +from airflow._shared.timezones import timezone +from airflow.models.base import COLLATION_ARGS, ID_LEN, Base +from airflow.utils.sqlalchemy import UtcDateTime + + +class TaskStateModel(Base): + """ + Persists key/value state for a task within a single DAG run. + + Scoped to (dag_run_id, task_id, map_index). Retries of the same task share + the same rows — that is the point. Different DAG runs have different dag_run_id + values so they get independent namespaces automatically. + """ + + __tablename__ = "task_state" + + dag_run_id: Mapped[int] = mapped_column(Integer, nullable=False, primary_key=True) + task_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False, primary_key=True) Review Comment: While here, switch the ORM to `StringID()` from `airflow.models.base` instead of the raw `String(ID_LEN, **COLLATION_ARGS)` form for `task_id`, `dag_id`, `run_id` (and the equivalents in `asset_state.py`). Every recent model (backfill, deadline, taskinstance) uses the helper -- keeps the collation in one place and matches what the migration should be doing per the previous comment. ########## airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py: ########## @@ -0,0 +1,84 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add task_state and asset_state tables. + +Revision ID: fde9ed84d07b +Revises: 9fabad868fdb +Create Date: 2026-04-27 09:29:20.165461 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import mysql + +from airflow.utils.sqlalchemy import UtcDateTime + +# revision identifiers, used by Alembic. +revision = "fde9ed84d07b" +down_revision = "9fabad868fdb" +branch_labels = None +depends_on = None +airflow_version = "3.3.0" + + +def upgrade(): + """Apply add task_state and asset_state tables.""" + op.create_table( + "asset_state", + sa.Column("asset_id", sa.Integer(), nullable=False), + sa.Column("key", sa.String(length=512), nullable=False), + sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(), "mysql"), nullable=False), + sa.Column("updated_at", UtcDateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["asset_id"], ["asset.id"], name="asset_state_asset_fkey", ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"), + ) + op.create_table( + "task_state", + sa.Column("dag_run_id", sa.Integer(), nullable=False), + sa.Column("task_id", sa.String(length=250), nullable=False), Review Comment: `task_id`, `dag_id`, `run_id` are declared as `sa.String(length=250)` here (and in the `asset_state` table), but the ORM uses `String(ID_LEN, **COLLATION_ARGS)` which on MySQL adds `utf8mb3_bin`. The migration ships with the server-default collation, so FK joins / WHERE comparisons across `task_state` and `dag_run` either error with `Illegal mix of collations` or skip the index on MySQL. Use `from airflow.migrations.db_types import StringID` and `StringID()` for these columns -- migration #0108 fixed the same drift in bulk for existing tables. -- 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]
