This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ebb01703096 AIP-103: Setting up foundation with ORM and Backend
Abstraction (#65759)
ebb01703096 is described below
commit ebb0170309695c4873249fa04d44bcc9fb5a5852
Author: Amogh Desai <[email protected]>
AuthorDate: Thu Apr 30 13:01:25 2026 +0530
AIP-103: Setting up foundation with ORM and Backend Abstraction (#65759)
* Adding ORM models for task_state and asset_state
* Adding base state backend and metastore implementation
* Introducing config: [state_store]
---
airflow-core/docs/migrations-ref.rst | 4 +-
airflow-core/pyproject.toml | 2 +
airflow-core/src/airflow/_shared/state | 1 +
.../src/airflow/config_templates/config.yml | 14 +
..._3_3_0_add_task_state_and_asset_state_tables.py | 85 +++++
airflow-core/src/airflow/models/__init__.py | 2 +
airflow-core/src/airflow/models/asset_state.py | 55 +++
airflow-core/src/airflow/models/task_state.py | 62 ++++
airflow-core/src/airflow/state/__init__.py | 38 ++
airflow-core/src/airflow/state/metastore.py | 359 ++++++++++++++++++
airflow-core/src/airflow/utils/db.py | 2 +-
airflow-core/tests/unit/models/test_asset_state.py | 150 ++++++++
airflow-core/tests/unit/models/test_task_state.py | 230 ++++++++++++
airflow-core/tests/unit/state/__init__.py | 18 +
airflow-core/tests/unit/state/test_metastore.py | 399 +++++++++++++++++++++
.../tests/test_pytest_args_for_test_types.py | 1 +
dev/breeze/tests/test_selective_checks.py | 40 +--
pyproject.toml | 5 +
scripts/ci/prek/check_tests_in_right_folders.py | 1 +
shared/state/.gitignore | 1 +
shared/state/.pre-commit-config.yaml | 31 ++
shared/state/pyproject.toml | 55 +++
shared/state/src/airflow_shared/state/__init__.py | 124 +++++++
shared/state/tests/conftest.py | 21 ++
shared/state/tests/state/__init__.py | 17 +
shared/state/tests/state/test_state.py | 72 ++++
uv.lock | 23 ++
27 files changed, 1790 insertions(+), 22 deletions(-)
diff --git a/airflow-core/docs/migrations-ref.rst
b/airflow-core/docs/migrations-ref.rst
index e54009b856e..55cbb7e4d1a 100644
--- a/airflow-core/docs/migrations-ref.rst
+++ b/airflow-core/docs/migrations-ref.rst
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are
executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description
|
+=========================+==================+===================+==============================================================+
-| ``9fabad868fdb`` (head) | ``a4c2d171ae18`` | ``3.3.0`` | Add
timetable_periodic to DagModel. |
+| ``fde9ed84d07b`` (head) | ``9fabad868fdb`` | ``3.3.0`` | Add
task_state and asset_state tables. |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``9fabad868fdb`` | ``a4c2d171ae18`` | ``3.3.0`` | Add
timetable_periodic to DagModel. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``a4c2d171ae18`` | ``1d6611b6ab7c`` | ``3.3.0`` | Add
dag_result to XComModel. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
diff --git a/airflow-core/pyproject.toml b/airflow-core/pyproject.toml
index b7f7f3c3b57..42d6dd737b6 100644
--- a/airflow-core/pyproject.toml
+++ b/airflow-core/pyproject.toml
@@ -247,6 +247,7 @@ exclude = [
"../shared/secrets_backend/src/airflow_shared/secrets_backend" =
"src/airflow/_shared/secrets_backend"
"../shared/secrets_masker/src/airflow_shared/secrets_masker" =
"src/airflow/_shared/secrets_masker"
"../shared/serialization/src/airflow_shared/serialization" =
"src/airflow/_shared/serialization"
+"../shared/state/src/airflow_shared/state" = "src/airflow/_shared/state"
"../shared/timezones/src/airflow_shared/timezones" =
"src/airflow/_shared/timezones"
"../shared/listeners/src/airflow_shared/listeners" =
"src/airflow/_shared/listeners"
"../shared/plugins_manager/src/airflow_shared/plugins_manager" =
"src/airflow/_shared/plugins_manager"
@@ -334,6 +335,7 @@ shared_distributions = [
"apache-airflow-shared-secrets-backend",
"apache-airflow-shared-secrets-masker",
"apache-airflow-shared-serialization",
+ "apache-airflow-shared-state",
"apache-airflow-shared-timezones",
"apache-airflow-shared-plugins-manager",
"apache-airflow-shared-providers-discovery",
diff --git a/airflow-core/src/airflow/_shared/state
b/airflow-core/src/airflow/_shared/state
new file mode 120000
index 00000000000..d251a2e1b0e
--- /dev/null
+++ b/airflow-core/src/airflow/_shared/state
@@ -0,0 +1 @@
+../../../../shared/state/src/airflow_shared/state
\ No newline at end of file
diff --git a/airflow-core/src/airflow/config_templates/config.yml
b/airflow-core/src/airflow/config_templates/config.yml
index c6dc3e97f29..3ac011a0999 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -2997,6 +2997,20 @@ dag_processor:
example: ~
default: "warning"
+state_store:
+ description: |
+ Configuration for task and asset state storage backend.
+ options:
+ backend:
+ description: |
+ Full dotted path to the class that implements state storage for tasks
and assets.
+ The class must be a subclass of ``BaseStateBackend``.
+ The default implementation persists state in the Airflow metadata
database.
+ version_added: 3.3.0
+ type: string
+ example: "mypackage.state.CustomStateBackend"
+ default: "airflow.state.metastore.MetastoreStateBackend"
+
profiling:
description: |
Configuration for memory profiling in Airflow component.
diff --git
a/airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py
b/airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py
new file mode 100644
index 00000000000..7f852d05c6c
--- /dev/null
+++
b/airflow-core/src/airflow/migrations/versions/0112_3_3_0_add_task_state_and_asset_state_tables.py
@@ -0,0 +1,85 @@
+#
+# 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.migrations.db_types import StringID
+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", StringID(), nullable=False),
+ sa.Column("map_index", sa.Integer(), server_default="-1",
nullable=False),
+ sa.Column("key", sa.String(length=512), nullable=False),
+ sa.Column("dag_id", StringID(), nullable=False),
+ sa.Column("run_id", StringID(), nullable=False),
+ sa.Column("value", sa.Text().with_variant(mysql.MEDIUMTEXT(),
"mysql"), nullable=False),
+ sa.Column("updated_at", UtcDateTime(), nullable=False),
+ sa.ForeignKeyConstraint(
+ ["dag_run_id"], ["dag_run.id"], name="task_state_dag_run_fkey",
ondelete="CASCADE"
+ ),
+ sa.PrimaryKeyConstraint("dag_run_id", "task_id", "map_index", "key",
name="task_state_pkey"),
+ )
+ with op.batch_alter_table("task_state", schema=None) as batch_op:
+ batch_op.create_index(
+ "idx_task_state_lookup", ["dag_id", "run_id", "task_id",
"map_index"], unique=False
+ )
+
+
+def downgrade():
+ """Unapply add task_state and asset_state tables."""
+ with op.batch_alter_table("task_state", schema=None) as batch_op:
+ batch_op.drop_index("idx_task_state_lookup")
+
+ op.drop_table("task_state")
+ op.drop_table("asset_state")
diff --git a/airflow-core/src/airflow/models/__init__.py
b/airflow-core/src/airflow/models/__init__.py
index 8e12325f568..9b134f4ff36 100644
--- a/airflow-core/src/airflow/models/__init__.py
+++ b/airflow-core/src/airflow/models/__init__.py
@@ -61,6 +61,7 @@ def import_all_models():
__getattr__(name)
import airflow.models.asset
+ import airflow.models.asset_state
import airflow.models.backfill
import airflow.models.dag_favorite
import airflow.models.dag_version
@@ -71,6 +72,7 @@ def import_all_models():
import airflow.models.errors
import airflow.models.revoked_token
import airflow.models.serialized_dag
+ import airflow.models.task_state
import airflow.models.taskinstancehistory
import airflow.models.tasklog
import airflow.models.team
diff --git a/airflow-core/src/airflow/models/asset_state.py
b/airflow-core/src/airflow/models/asset_state.py
new file mode 100644
index 00000000000..ecae12acaa8
--- /dev/null
+++ b/airflow-core/src/airflow/models/asset_state.py
@@ -0,0 +1,55 @@
+#
+# 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, 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, Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class AssetStateModel(Base):
+ """
+ Persists key/value state scoped to an asset identity.
+
+ Not scoped to any DAG run — a watermark written in run 1 is readable by
run 2.
+ Rows survive until explicitly deleted or the asset itself is deleted.
+ """
+
+ __tablename__ = "asset_state"
+
+ asset_id: Mapped[int] = mapped_column(Integer, nullable=False,
primary_key=True)
+ key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS),
nullable=False, primary_key=True)
+
+ value: Mapped[str] = mapped_column(Text().with_variant(MEDIUMTEXT,
"mysql"), nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
+
+ __table_args__ = (
+ PrimaryKeyConstraint("asset_id", "key", name="asset_state_pkey"),
+ ForeignKeyConstraint(
+ ["asset_id"],
+ ["asset.id"],
+ name="asset_state_asset_fkey",
+ ondelete="CASCADE",
+ ),
+ )
diff --git a/airflow-core/src/airflow/models/task_state.py
b/airflow-core/src/airflow/models/task_state.py
new file mode 100644
index 00000000000..dbc17e3b069
--- /dev/null
+++ b/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, Base, StringID
+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(StringID(), nullable=False,
primary_key=True)
+ map_index: Mapped[int] = mapped_column(Integer, primary_key=True,
nullable=False, server_default="-1")
+ key: Mapped[str] = mapped_column(String(512, **COLLATION_ARGS),
nullable=False, primary_key=True)
+
+ dag_id: Mapped[str] = mapped_column(StringID(), nullable=False)
+ run_id: Mapped[str] = mapped_column(StringID(), nullable=False)
+
+ value: Mapped[str] = mapped_column(Text().with_variant(MEDIUMTEXT,
"mysql"), nullable=False)
+ updated_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
+
+ __table_args__ = (
+ PrimaryKeyConstraint("dag_run_id", "task_id", "map_index", "key",
name="task_state_pkey"),
+ ForeignKeyConstraint(
+ ["dag_run_id"],
+ ["dag_run.id"],
+ name="task_state_dag_run_fkey",
+ ondelete="CASCADE",
+ ),
+ Index("idx_task_state_lookup", "dag_id", "run_id", "task_id",
"map_index"),
+ )
diff --git a/airflow-core/src/airflow/state/__init__.py
b/airflow-core/src/airflow/state/__init__.py
new file mode 100644
index 00000000000..89fa0801a06
--- /dev/null
+++ b/airflow-core/src/airflow/state/__init__.py
@@ -0,0 +1,38 @@
+#
+# 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 airflow._shared.state import (
+ AssetScope as AssetScope,
+ BaseStateBackend as BaseStateBackend,
+ StateScope as StateScope,
+ TaskScope as TaskScope,
+)
+
+
+def resolve_state_backend() -> type[BaseStateBackend]:
+ from airflow.configuration import conf
+
+ clazz = conf.getimport("state_store", "backend")
+ if clazz is None:
+ raise ValueError("state_store.backend is not configured or resolved to
None.")
+ if not issubclass(clazz, BaseStateBackend):
+ raise TypeError(
+ f"Your custom state backend `{clazz.__name__}` is not a subclass
of `BaseStateBackend`."
+ )
+ return clazz
diff --git a/airflow-core/src/airflow/state/metastore.py
b/airflow-core/src/airflow/state/metastore.py
new file mode 100644
index 00000000000..3382dad81fc
--- /dev/null
+++ b/airflow-core/src/airflow/state/metastore.py
@@ -0,0 +1,359 @@
+#
+# 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.typing_compat import assert_never
+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)
+ case _:
+ assert_never(scope)
+
+ @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)
+ case _:
+ assert_never(scope)
+
+ @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)
+ case _:
+ assert_never(scope)
+
+ @provide_session
+ def clear(
+ self,
+ scope: StateScope,
+ *,
+ all_map_indices: bool = False,
+ session: Session = NEW_SESSION,
+ ) -> None:
+ match scope:
+ case TaskScope():
+ self._clear_task_state(scope, all_map_indices=all_map_indices,
session=session)
+ case AssetScope():
+ self._clear_asset_state(scope, session=session)
+ case _:
+ assert_never(scope)
+
+ 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)
+ case _:
+ assert_never(scope)
+
+ 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)
+ case _:
+ assert_never(scope)
+
+ 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)
+ case _:
+ assert_never(scope)
+
+ async def aclear(self, scope: StateScope, *, all_map_indices: bool =
False) -> None:
+ async with create_session_async() as session:
+ match scope:
+ case TaskScope():
+ await self._aclear_task_state(scope,
all_map_indices=all_map_indices, session=session)
+ case AssetScope():
+ await self._aclear_asset_state(scope, session=session)
+ case _:
+ assert_never(scope)
+
+ 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, *, all_map_indices: bool =
False, session: Session) -> None:
+ conditions = [
+ TaskStateModel.dag_id == scope.dag_id,
+ TaskStateModel.run_id == scope.run_id,
+ TaskStateModel.task_id == scope.task_id,
+ ]
+ if not all_map_indices:
+ conditions.append(TaskStateModel.map_index == scope.map_index)
+ session.execute(delete(TaskStateModel).where(*conditions))
+
+ def _get_asset_state(self, scope: AssetScope, key: str, *, session:
Session) -> str | None:
+ row = session.scalar(
+ select(AssetStateModel).where(
+ AssetStateModel.asset_id == scope.asset_id,
+ AssetStateModel.key == key,
+ )
+ )
+ return row.value if row is not None else None
+
+ def _set_asset_state(self, scope: AssetScope, key: str, value: str, *,
session: Session) -> None:
+ now = timezone.utcnow()
+ values = dict(asset_id=scope.asset_id, key=key, value=value,
updated_at=now)
+ stmt = _build_upsert_stmt(
+ get_dialect_name(session),
+ AssetStateModel,
+ ["asset_id", "key"],
+ values,
+ dict(value=value, updated_at=now),
+ )
+ session.execute(stmt)
+
+ def _delete_asset_state(self, scope: AssetScope, key: str, *, session:
Session) -> None:
+ session.execute(
+ delete(AssetStateModel).where(
+ AssetStateModel.asset_id == scope.asset_id,
+ AssetStateModel.key == key,
+ )
+ )
+
+ def _clear_asset_state(self, scope: AssetScope, *, session: Session) ->
None:
+ session.execute(
+ delete(AssetStateModel).where(
+ AssetStateModel.asset_id == scope.asset_id,
+ )
+ )
+
+ async def _aget_task_state(self, scope: TaskScope, key: str, *, session:
AsyncSession) -> str | None:
+ row = await 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
+
+ async def _aset_task_state(
+ self, scope: TaskScope, key: str, value: str, *, session: AsyncSession
+ ) -> None:
+ dag_run_id = await 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,
+ )
+ # get_dialect_name expects a sync Session; sync_session is the
underlying Session the async wrapper delegates to
+ stmt = _build_upsert_stmt(
+ get_dialect_name(session.sync_session),
+ TaskStateModel,
+ ["dag_run_id", "task_id", "map_index", "key"],
+ values,
+ dict(value=value, updated_at=now),
+ )
+ await session.execute(stmt)
+
+ async def _adelete_task_state(self, scope: TaskScope, key: str, *,
session: AsyncSession) -> None:
+ await 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,
+ )
+ )
+
+ async def _aclear_task_state(
+ self, scope: TaskScope, *, all_map_indices: bool = False, session:
AsyncSession
+ ) -> None:
+ conditions = [
+ TaskStateModel.dag_id == scope.dag_id,
+ TaskStateModel.run_id == scope.run_id,
+ TaskStateModel.task_id == scope.task_id,
+ ]
+ if not all_map_indices:
+ conditions.append(TaskStateModel.map_index == scope.map_index)
+ await session.execute(delete(TaskStateModel).where(*conditions))
+
+ async def _aget_asset_state(self, scope: AssetScope, key: str, *, session:
AsyncSession) -> str | None:
+ row = await session.scalar(
+ select(AssetStateModel).where(
+ AssetStateModel.asset_id == scope.asset_id,
+ AssetStateModel.key == key,
+ )
+ )
+ return row.value if row is not None else None
+
+ async def _aset_asset_state(
+ self, scope: AssetScope, key: str, value: str, *, session: AsyncSession
+ ) -> None:
+ now = timezone.utcnow()
+ values = dict(asset_id=scope.asset_id, key=key, value=value,
updated_at=now)
+ # get_dialect_name expects a sync Session; sync_session is the
underlying Session the async wrapper delegates to
+ stmt = _build_upsert_stmt(
+ get_dialect_name(session.sync_session),
+ AssetStateModel,
+ ["asset_id", "key"],
+ values,
+ dict(value=value, updated_at=now),
+ )
+ await session.execute(stmt)
+
+ async def _adelete_asset_state(self, scope: AssetScope, key: str, *,
session: AsyncSession) -> None:
+ await session.execute(
+ delete(AssetStateModel).where(
+ AssetStateModel.asset_id == scope.asset_id,
+ AssetStateModel.key == key,
+ )
+ )
+
+ async def _aclear_asset_state(self, scope: AssetScope, *, session:
AsyncSession) -> None:
+ await session.execute(
+ delete(AssetStateModel).where(
+ AssetStateModel.asset_id == scope.asset_id,
+ )
+ )
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index 9293f1c49ac..0e1173c2dce 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -116,7 +116,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.1.0": "cc92b33c6709",
"3.1.8": "509b94a1042d",
"3.2.0": "1d6611b6ab7c",
- "3.3.0": "9fabad868fdb",
+ "3.3.0": "fde9ed84d07b",
}
# Prefix used to identify tables holding data moved during migration.
diff --git a/airflow-core/tests/unit/models/test_asset_state.py
b/airflow-core/tests/unit/models/test_asset_state.py
new file mode 100644
index 00000000000..cb2976fe411
--- /dev/null
+++ b/airflow-core/tests/unit/models/test_asset_state.py
@@ -0,0 +1,150 @@
+#
+# 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 delete, select
+from sqlalchemy.exc import IntegrityError
+
+from airflow.models.asset import AssetModel
+from airflow.models.asset_state import AssetStateModel
+
+from tests_common.test_utils.db import clear_db_assets
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
+pytestmark = pytest.mark.db_test
+
+
[email protected](autouse=True)
+def clean_tables():
+ clear_db_assets()
+ yield
+ clear_db_assets()
+
+
[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
+
+
+class TestAssetStateModel:
+ def test_insert_and_read(self, session: Session, asset: AssetModel):
+ """A row written for an asset key is readable back with the correct
value."""
+ row = AssetStateModel(
+ asset_id=asset.id,
+ key="last_processed_at",
+ value="2026-04-24T00:00:00Z",
+ )
+ session.add(row)
+ session.flush()
+
+ result = session.scalar(
+ select(AssetStateModel).where(
+ AssetStateModel.asset_id == asset.id,
+ AssetStateModel.key == "last_processed_at",
+ )
+ )
+ assert result is not None
+ assert result.value == "2026-04-24T00:00:00Z"
+
+ def test_duplicate_pk_raises(self, session: Session, asset: AssetModel):
+ """Inserting a second row with the same (asset_id, key) raises
IntegrityError."""
+ session.add(AssetStateModel(asset_id=asset.id, key="watermark",
value="first"))
+ session.flush()
+
+ session.add(AssetStateModel(asset_id=asset.id, key="watermark",
value="duplicate"))
+ with pytest.raises(IntegrityError):
+ session.flush()
+
+ def test_multiple_keys_per_asset(self, session: Session, asset:
AssetModel):
+ """An asset can hold multiple independent keys in the same
namespace."""
+ for key in ("watermark", "last_file_count", "last_error"):
+ session.add(AssetStateModel(asset_id=asset.id, key=key,
value=f"val_{key}"))
+ session.flush()
+
+ results =
session.scalars(select(AssetStateModel).where(AssetStateModel.asset_id ==
asset.id)).all()
+ assert len(results) == 3
+
+ def test_cascade_delete_on_asset(self, session: Session, asset:
AssetModel):
+ """Deleting an asset cascades to its asset_state rows — no orphans
remain."""
+ session.add(
+ AssetStateModel(
+ asset_id=asset.id,
+ key="last_processed_at",
+ value="2026-04-24T00:00:00Z",
+ )
+ )
+ session.flush()
+
+ session.execute(delete(AssetModel).where(AssetModel.id == asset.id))
+ session.flush()
+
+ remaining =
session.scalars(select(AssetStateModel).where(AssetStateModel.asset_id ==
asset.id)).all()
+ assert remaining == []
+
+ def test_different_assets_are_isolated(self, session: Session, asset:
AssetModel):
+ """State written for one asset is not visible when querying a
different asset."""
+ asset2 = AssetModel(uri="s3://bucket/other", name="other_asset",
group="test")
+ session.add(asset2)
+ session.flush()
+
+ session.add(
+ AssetStateModel(
+ asset_id=asset.id,
+ key="watermark",
+ value="asset1_value",
+ )
+ )
+ session.flush()
+
+ result = session.scalar(
+ select(AssetStateModel).where(
+ AssetStateModel.asset_id == asset2.id,
+ AssetStateModel.key == "watermark",
+ )
+ )
+ assert result is None
+
+ def test_state_survives_across_runs(self, session: Session, asset:
AssetModel):
+ """Asset state is not scoped to a DAG run — updating a value persists
it for future reads."""
+ row = AssetStateModel(
+ asset_id=asset.id,
+ key="watermark",
+ value="2026-04-01T00:00:00Z",
+ )
+ session.add(row)
+ session.flush()
+
+ row.value = "2026-04-24T00:00:00Z"
+ session.flush()
+
+ result = session.scalar(
+ select(AssetStateModel).where(
+ AssetStateModel.asset_id == asset.id,
+ AssetStateModel.key == "watermark",
+ )
+ )
+ assert result is not None
+ assert result.value == "2026-04-24T00:00:00Z"
diff --git a/airflow-core/tests/unit/models/test_task_state.py
b/airflow-core/tests/unit/models/test_task_state.py
new file mode 100644
index 00000000000..5e527ec8ebd
--- /dev/null
+++ b/airflow-core/tests/unit/models/test_task_state.py
@@ -0,0 +1,230 @@
+#
+# 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 delete, select
+from sqlalchemy.exc import IntegrityError
+
+from airflow._shared.timezones import timezone
+from airflow.models.dagrun import DagRun, DagRunType
+from airflow.models.task_state import TaskStateModel
+
+from tests_common.test_utils.db import 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()
+ yield
+ clear_db_dags()
+ clear_db_runs()
+
+
[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
+
+
+class TestTaskStateModel:
+ def test_insert_and_read(self, session: Session, dag_run: DagRun):
+ row = TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=-1,
+ key="remote_job_id",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value="application_1234",
+ )
+ session.add(row)
+ session.flush()
+
+ result = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_run_id == dag_run.id,
+ TaskStateModel.task_id == TASK_ID,
+ TaskStateModel.map_index == -1,
+ TaskStateModel.key == "remote_job_id",
+ )
+ )
+ assert result is not None
+ assert result.value == "application_1234"
+ assert result.dag_id == DAG_ID
+ assert result.run_id == RUN_ID
+
+ def test_duplicate_pk_raises(self, session: Session, dag_run: DagRun):
+ """Inserting a second row with the same (dag_run_id, task_id,
map_index, key) raises IntegrityError."""
+ session.add(
+ TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=-1,
+ key="remote_job_id",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value="first",
+ )
+ )
+ session.flush()
+
+ session.add(
+ TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=-1,
+ key="remote_job_id",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value="duplicate",
+ )
+ )
+ with pytest.raises(IntegrityError):
+ session.flush()
+
+ def test_retries_share_state(self, session: Session, dag_run: DagRun):
+ """Retries of the same task share the same row — updates are visible
across retries."""
+ row = TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=-1,
+ key="remote_job_id",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value="application_try1",
+ )
+ session.add(row)
+ session.flush()
+
+ row.value = "application_try2"
+ session.flush()
+
+ result = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_run_id == dag_run.id,
+ TaskStateModel.task_id == TASK_ID,
+ TaskStateModel.key == "remote_job_id",
+ )
+ )
+ assert result is not None
+ assert result.value == "application_try2"
+
+ def test_map_index_isolation(self, session: Session, dag_run: DagRun):
+ """Each mapped task instance gets its own independent state
namespace."""
+ for idx in (0, 1, 2):
+ session.add(
+ TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=idx,
+ key="checkpoint",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value=f"value_{idx}",
+ )
+ )
+ session.flush()
+
+ for idx in (0, 1, 2):
+ result = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_run_id == dag_run.id,
+ TaskStateModel.task_id == TASK_ID,
+ TaskStateModel.map_index == idx,
+ TaskStateModel.key == "checkpoint",
+ )
+ )
+ assert result is not None
+ assert result.value == f"value_{idx}"
+
+ def test_cascade_delete_on_dag_run(self, session: Session, dag_run:
DagRun):
+ """Deleting a DagRun cascades to its task_state rows — no orphans
remain."""
+ session.add(
+ TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=-1,
+ key="remote_job_id",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value="application_1234",
+ )
+ )
+ session.flush()
+
+ session.execute(delete(DagRun).where(DagRun.id == dag_run.id))
+ session.flush()
+
+ remaining = session.scalars(
+ select(TaskStateModel).where(TaskStateModel.dag_run_id ==
dag_run.id)
+ ).all()
+ assert remaining == []
+
+ def test_different_dag_runs_are_isolated(self, session: Session, dag_run:
DagRun):
+ """State written for one DagRun is not visible when querying a
different DagRun."""
+ run2 = DagRun(
+ dag_id=DAG_ID,
+ run_id="scheduled__2026-04-25",
+ run_type=DagRunType.SCHEDULED,
+ logical_date=timezone.datetime(2026, 4, 25),
+ run_after=timezone.datetime(2026, 4, 25),
+ )
+ session.add(run2)
+ session.flush()
+
+ session.add(
+ TaskStateModel(
+ dag_run_id=dag_run.id,
+ task_id=TASK_ID,
+ map_index=-1,
+ key="watermark",
+ dag_id=DAG_ID,
+ run_id=RUN_ID,
+ value="run1_value",
+ )
+ )
+ session.flush()
+
+ result = session.scalar(
+ select(TaskStateModel).where(
+ TaskStateModel.dag_run_id == run2.id,
+ TaskStateModel.key == "watermark",
+ )
+ )
+ assert result is None
diff --git a/airflow-core/tests/unit/state/__init__.py
b/airflow-core/tests/unit/state/__init__.py
new file mode 100644
index 00000000000..03cb33c14c4
--- /dev/null
+++ b/airflow-core/tests/unit/state/__init__.py
@@ -0,0 +1,18 @@
+#
+# 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
diff --git a/airflow-core/tests/unit/state/test_metastore.py
b/airflow-core/tests/unit/state/test_metastore.py
new file mode 100644
index 00000000000..98993d7133c
--- /dev/null
+++ b/airflow-core/tests/unit/state/test_metastore.py
@@ -0,0 +1,399 @@
+#
+# 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"
+
+ @pytest.mark.backend("postgres", "mysql", "sqlite")
+ 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"
+
+ def test_clear_with_all_map_indices_flag_wipes_wide(
+ self, session: Session, backend: MetastoreStateBackend, dag_run: DagRun
+ ):
+ """clear(scope, all_map_indices=True) removes state for every map
index of the task."""
+ 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, all_map_indices=True, session=session)
+ session.flush()
+
+ assert backend.get(scope0, "job_id", session=session) is None
+ assert backend.get(scope1, "job_id", session=session) is None
+
+
+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"
+
+ @pytest.mark.backend("postgres", "mysql", "sqlite")
+ 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_aclear_with_all_map_indices_flag_wipes_fleet(
+ self, backend: MetastoreStateBackend, dag_run_committed: DagRun
+ ):
+ """aclear(scope, all_map_indices=True) removes state for every map
index of the task."""
+ 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)
+
+ await backend.aset(scope0, "job_id", "app_0")
+ await backend.aset(scope1, "job_id", "app_1")
+
+ await backend.aclear(scope0, all_map_indices=True)
+
+ assert await backend.aget(scope0, "job_id") is None
+ assert await backend.aget(scope1, "job_id") 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:
+ @conf_vars({("state_store", "backend"):
"airflow.state.metastore.MetastoreStateBackend"})
+ def test_resolve_returns_configured_backend(self):
+ """resolve_state_backend() imports and returns the explicitly
configured backend class."""
+ assert resolve_state_backend() is MetastoreStateBackend
+
+ @conf_vars({("state_store", "backend"): ""})
+ def test_empty_backend_raises_value_error(self):
+ """resolve_state_backend() raises ValueError when backend is
explicitly set to empty string."""
+ with pytest.raises(ValueError, match="state_store.backend is not
configured"):
+ resolve_state_backend()
+
+ @conf_vars({("state_store", "backend"): "airflow.models.dagrun.DagRun"})
+ def test_invalid_backend_raises_type_error(self):
+ """resolve_state_backend() raises TypeError when the configured class
is not a BaseStateBackend subclass."""
+ with pytest.raises(TypeError, match="not a subclass of
`BaseStateBackend`"):
+ resolve_state_backend()
diff --git a/dev/breeze/tests/test_pytest_args_for_test_types.py
b/dev/breeze/tests/test_pytest_args_for_test_types.py
index 423c96eb10f..e184a70a357 100644
--- a/dev/breeze/tests/test_pytest_args_for_test_types.py
+++ b/dev/breeze/tests/test_pytest_args_for_test_types.py
@@ -173,6 +173,7 @@ def _find_all_integration_folders() -> list[str]:
"airflow-core/tests/unit/plugins",
"airflow-core/tests/unit/security",
"airflow-core/tests/unit/sensors",
+ "airflow-core/tests/unit/state",
"airflow-core/tests/unit/task",
"airflow-core/tests/unit/testconfig",
"airflow-core/tests/unit/timetables",
diff --git a/dev/breeze/tests/test_selective_checks.py
b/dev/breeze/tests/test_selective_checks.py
index 2922465ca66..7ffd1370475 100644
--- a/dev/breeze/tests/test_selective_checks.py
+++ b/dev/breeze/tests/test_selective_checks.py
@@ -107,7 +107,7 @@ ALL_SKIPPED_COMMITS_ON_NO_CI_IMAGE = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
@@ -120,7 +120,7 @@ ALL_SKIPPED_COMMITS_IF_NO_UI = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
ALL_SKIPPED_COMMITS_IF_NO_HELM_TESTS = (
@@ -131,7 +131,7 @@ ALL_SKIPPED_COMMITS_IF_NO_HELM_TESTS = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,update-uv-lock"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,update-uv-lock"
)
ALL_SKIPPED_COMMITS_IF_NO_UI_AND_HELM_TESTS = (
@@ -142,7 +142,7 @@ ALL_SKIPPED_COMMITS_IF_NO_UI_AND_HELM_TESTS = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
@@ -154,7 +154,7 @@ ALL_SKIPPED_COMMITS_IF_NO_PROVIDERS_AND_UI = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
@@ -166,7 +166,7 @@ ALL_SKIPPED_COMMITS_IF_NO_PROVIDERS = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
@@ -179,7 +179,7 @@ ALL_SKIPPED_COMMITS_IF_NO_PROVIDERS_UI_AND_HELM_TESTS = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
@@ -191,7 +191,7 @@ ALL_SKIPPED_COMMITS_IF_NO_CODE_PROVIDERS_AND_HELM_TESTS = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,update-uv-lock"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,update-uv-lock"
)
ALL_SKIPPED_COMMITS_IF_NOT_IMPORTANT_FILES_CHANGED = (
@@ -202,7 +202,7 @@ ALL_SKIPPED_COMMITS_IF_NOT_IMPORTANT_FILES_CHANGED = (
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
)
@@ -405,7 +405,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str,
str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"upgrade-to-newer-dependencies": "false",
@@ -678,7 +678,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str,
str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"skip-providers-tests": "false",
@@ -716,7 +716,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str,
str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"skip-providers-tests": "true",
@@ -753,7 +753,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str,
str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"skip-providers-tests": "true",
@@ -789,7 +789,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str,
str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"skip-providers-tests": "true",
@@ -1130,7 +1130,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"upgrade-to-newer-dependencies": "false",
@@ -1294,7 +1294,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"upgrade-to-newer-dependencies": "false",
@@ -1454,7 +1454,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
"mypy-shared-module_loading,mypy-shared-observability,"
"mypy-shared-plugins_manager,mypy-shared-providers_discovery,"
"mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,"
"mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
@@ -2191,7 +2191,7 @@ def test_expected_output_push(
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"upgrade-to-newer-dependencies": "false",
@@ -2232,7 +2232,7 @@ def test_expected_output_push(
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"run-kubernetes-tests": "true",
@@ -2278,7 +2278,7 @@ def test_expected_output_push(
"mypy-shared-configuration,mypy-shared-dagnode,mypy-shared-listeners,mypy-shared-logging,"
"mypy-shared-module_loading,mypy-shared-observability,mypy-shared-plugins_manager,"
"mypy-shared-providers_discovery,mypy-shared-secrets_backend,mypy-shared-secrets_masker,"
-
"mypy-shared-serialization,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
+
"mypy-shared-serialization,mypy-shared-state,mypy-shared-template_rendering,mypy-shared-timezones,mypy-task-sdk,mypy-task-sdk-integration-tests,"
"ts-compile-lint-simple-auth-manager-ui,ts-compile-lint-ui,update-uv-lock"
),
"run-kubernetes-tests": "false",
diff --git a/pyproject.toml b/pyproject.toml
index bf8bd0f272e..e085c79b998 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1326,6 +1326,7 @@ dev = [
"apache-airflow-shared-secrets-backend",
"apache-airflow-shared-secrets-masker",
"apache-airflow-shared-serialization",
+ "apache-airflow-shared-state",
"apache-airflow-shared-template-rendering",
"apache-airflow-shared-timezones",
]
@@ -1492,6 +1493,7 @@ apache-airflow-shared-providers-discovery = false
apache-airflow-shared-secrets-backend = false
apache-airflow-shared-secrets-masker = false
apache-airflow-shared-serialization = false
+apache-airflow-shared-state = false
apache-airflow-shared-template-rendering = false
apache-airflow-shared-timezones = false
apache-airflow-task-sdk = false
@@ -1642,6 +1644,7 @@ apache-airflow-shared-providers-discovery = false
apache-airflow-shared-secrets-backend = false
apache-airflow-shared-secrets-masker = false
apache-airflow-shared-serialization = false
+apache-airflow-shared-state = false
apache-airflow-shared-template-rendering = false
apache-airflow-shared-timezones = false
apache-airflow-task-sdk = false
@@ -1686,6 +1689,7 @@ apache-airflow-shared-providers-discovery = { workspace =
true }
apache-airflow-shared-secrets-backend = { workspace = true }
apache-airflow-shared-secrets-masker = { workspace = true }
apache-airflow-shared-serialization = { workspace = true }
+apache-airflow-shared-state = { workspace = true }
apache-airflow-shared-template-rendering = { workspace = true }
apache-airflow-shared-timezones = { workspace = true }
# Automatically generated provider workspace items
(update_airflow_pyproject_toml.py)
@@ -1822,6 +1826,7 @@ members = [
"shared/secrets_backend",
"shared/secrets_masker",
"shared/serialization",
+ "shared/state",
"shared/template_rendering",
"shared/timezones",
# Automatically generated provider workspace members
(update_airflow_pyproject_toml.py)
diff --git a/scripts/ci/prek/check_tests_in_right_folders.py
b/scripts/ci/prek/check_tests_in_right_folders.py
index 1a2bf5b53bc..7c164c0bb14 100755
--- a/scripts/ci/prek/check_tests_in_right_folders.py
+++ b/scripts/ci/prek/check_tests_in_right_folders.py
@@ -72,6 +72,7 @@ POSSIBLE_TEST_FOLDERS = [
"security",
"sensors",
"serialization",
+ "state",
"task",
"template",
"test_utils",
diff --git a/shared/state/.gitignore b/shared/state/.gitignore
new file mode 100644
index 00000000000..bff2d762960
--- /dev/null
+++ b/shared/state/.gitignore
@@ -0,0 +1 @@
+*.iml
diff --git a/shared/state/.pre-commit-config.yaml
b/shared/state/.pre-commit-config.yaml
new file mode 100644
index 00000000000..9c22fb9a0ec
--- /dev/null
+++ b/shared/state/.pre-commit-config.yaml
@@ -0,0 +1,31 @@
+# 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.
+---
+default_stages: [pre-commit, pre-push]
+minimum_prek_version: '0.3.4'
+default_language_version:
+ python: python3
+repos:
+ - repo: local
+ hooks:
+ - id: mypy-shared-state
+ name: Run mypy for shared-state
+ language: python
+ entry:
../../scripts/ci/prek/run_mypy_full_dist_local_venv_or_breeze_in_ci.py
shared/state
+ pass_filenames: false
+ files: ^.*\.py$
+ require_serial: true
diff --git a/shared/state/pyproject.toml b/shared/state/pyproject.toml
new file mode 100644
index 00000000000..f16ea66797a
--- /dev/null
+++ b/shared/state/pyproject.toml
@@ -0,0 +1,55 @@
+# 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.
+
+[project]
+name = "apache-airflow-shared-state"
+description = "Shared task and asset state abstractions for Airflow
distributions"
+version = "0.0"
+classifiers = [
+ "Private :: Do Not Upload",
+]
+
+dependencies = []
+
+[dependency-groups]
+dev = [
+ "apache-airflow-devel-common",
+]
+mypy = [
+ "apache-airflow-devel-common[mypy]",
+]
+
+[build-system]
+requires = [
+ "hatchling==1.29.0",
+ "packaging==26.1",
+ "pathspec==1.0.4",
+ "pluggy==1.6.0",
+ "tomli==2.4.1; python_version < '3.11'",
+ "trove-classifiers==2026.1.14.14",
+]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/airflow_shared"]
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+src = ["src"]
+
+[tool.ruff.lint.per-file-ignores]
+"!src/*" = ["D", "S101", "TRY002"]
diff --git a/shared/state/src/airflow_shared/state/__init__.py
b/shared/state/src/airflow_shared/state/__init__.py
new file mode 100644
index 00000000000..463d9f378f3
--- /dev/null
+++ b/shared/state/src/airflow_shared/state/__init__.py
@@ -0,0 +1,124 @@
+# 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 abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TaskScope:
+ """
+ Identifies the state namespace for a single task instance (or retry
thereof).
+
+ ``map_index`` defaults to ``-1`` for non-mapped tasks. For mapped tasks,
+ set it to the actual mapped index. ``get``/``set``/``delete`` always match
+ on ``(dag_id, run_id, task_id, map_index)`` exactly. To wipe state across
+ every map index of the task, call ``clear``/``aclear`` with
+ ``all_map_indices=True``.
+ """
+
+ dag_id: str
+ run_id: str
+ task_id: str
+ map_index: int = -1
+
+
+@dataclass(frozen=True)
+class AssetScope:
+ """Identifies the state namespace for an asset."""
+
+ asset_id: int
+
+
+StateScope = TaskScope | AssetScope
+
+
+class BaseStateBackend(ABC):
+ """
+ Abstract backend for reading and writing task and asset state.
+
+ Each method receives a ``StateScope`` which is either a ``TaskScope`` or
an ``AssetScope``.
+ Implementations must handle both types. The standard dispatch pattern is::
+
+ match scope:
+ case TaskScope():
+ ... # task-specific storage
+ case AssetScope():
+ ... # asset-specific storage
+
+ Custom backends are configured via ``[state_store] backend`` in
``airflow.cfg``.
+ """
+
+ @abstractmethod
+ def get(self, scope: StateScope, key: str) -> 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:
+ """
+ Write or overwrite the value for the given key.
+
+ Must handle both ``TaskScope`` and ``AssetScope``.
+ """
+
+ @abstractmethod
+ def delete(self, scope: StateScope, key: str) -> 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:
+ """
+ Delete all keys under the given scope.
+
+ 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 drop the
``map_index``
+ filter entirely and wipe state across every mapped instance of the
task.
+ For ``AssetScope`` the flag has no effect.
+ """
+
+ @abstractmethod
+ async def aget(self, scope: StateScope, key: str) -> str | None:
+ """Async variant of get. Must handle both ``TaskScope`` and
``AssetScope``."""
+
+ @abstractmethod
+ async def aset(self, scope: StateScope, key: str, value: str) -> None:
+ """Async variant of set. Must handle both ``TaskScope`` and
``AssetScope``."""
+
+ @abstractmethod
+ async def adelete(self, scope: StateScope, key: str) -> None:
+ """Async variant of delete. Must handle both ``TaskScope`` and
``AssetScope``."""
+
+ @abstractmethod
+ async def aclear(self, scope: StateScope, *, all_map_indices: bool =
False) -> 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.
+ """
diff --git a/shared/state/tests/conftest.py b/shared/state/tests/conftest.py
new file mode 100644
index 00000000000..93aecf26184
--- /dev/null
+++ b/shared/state/tests/conftest.py
@@ -0,0 +1,21 @@
+# 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 os
+
+os.environ["_AIRFLOW__AS_LIBRARY"] = "true"
diff --git a/shared/state/tests/state/__init__.py
b/shared/state/tests/state/__init__.py
new file mode 100644
index 00000000000..21d298ede6e
--- /dev/null
+++ b/shared/state/tests/state/__init__.py
@@ -0,0 +1,17 @@
+# 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
diff --git a/shared/state/tests/state/test_state.py
b/shared/state/tests/state/test_state.py
new file mode 100644
index 00000000000..47bce18a69e
--- /dev/null
+++ b/shared/state/tests/state/test_state.py
@@ -0,0 +1,72 @@
+# 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 pytest
+
+from airflow_shared.state import BaseStateBackend, StateScope
+
+
+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, *, all_map_indices: bool =
False) -> 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, *, all_map_indices: bool
= False) -> None:
+ self.calls.append("aclear")
+
+ return ConcreteBackend()
+
+ def test_incomplete_subclass_raises_type_error(self):
+ """A subclass that omits any abstract method cannot be instantiated."""
+
+ class IncompleteBackend(BaseStateBackend):
+ def get(self, scope: StateScope, key: str) -> str | None:
+ return None
+
+ with pytest.raises(TypeError, match="Can't instantiate abstract
class"):
+ IncompleteBackend()
+
+ def test_abstract_methods_cover_full_interface(self):
+ """BaseStateBackend enforces all 8 sync+async methods as abstract."""
+ expected = {"get", "set", "delete", "clear", "aget", "aset",
"adelete", "aclear"}
+ assert BaseStateBackend.__abstractmethods__ == expected
diff --git a/uv.lock b/uv.lock
index 0c5a314f249..73aa0f11a7d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -89,6 +89,7 @@ apache-airflow-providers-http = false
apache-airflow-providers-slack = false
apache-airflow-providers-vespa = false
apache-airflow-providers-databricks = false
+apache-airflow-shared-state = false
apache-airflow-providers-sqlite = false
apache-airflow-shared-module-loading = false
apache-airflow-providers-yandex = false
@@ -281,6 +282,7 @@ members = [
"apache-airflow-shared-secrets-backend",
"apache-airflow-shared-secrets-masker",
"apache-airflow-shared-serialization",
+ "apache-airflow-shared-state",
"apache-airflow-shared-template-rendering",
"apache-airflow-shared-timezones",
"apache-airflow-task-sdk",
@@ -1435,6 +1437,7 @@ dev = [
{ name = "apache-airflow-shared-secrets-backend" },
{ name = "apache-airflow-shared-secrets-masker" },
{ name = "apache-airflow-shared-serialization" },
+ { name = "apache-airflow-shared-state" },
{ name = "apache-airflow-shared-template-rendering" },
{ name = "apache-airflow-shared-timezones" },
{ name = "apache-airflow-task-sdk", extra = ["all"] },
@@ -1707,6 +1710,7 @@ dev = [
{ name = "apache-airflow-shared-secrets-backend", editable =
"shared/secrets_backend" },
{ name = "apache-airflow-shared-secrets-masker", editable =
"shared/secrets_masker" },
{ name = "apache-airflow-shared-serialization", editable =
"shared/serialization" },
+ { name = "apache-airflow-shared-state", editable = "shared/state" },
{ name = "apache-airflow-shared-template-rendering", editable =
"shared/template_rendering" },
{ name = "apache-airflow-shared-timezones", editable = "shared/timezones"
},
{ name = "apache-airflow-task-sdk", extras = ["all"], editable =
"task-sdk" },
@@ -8360,6 +8364,25 @@ mypy = [
dev = [{ name = "apache-airflow-devel-common", editable = "devel-common" }]
mypy = [{ name = "apache-airflow-devel-common", extras = ["mypy"], editable =
"devel-common" }]
+[[package]]
+name = "apache-airflow-shared-state"
+version = "0.0"
+source = { editable = "shared/state" }
+
+[package.dev-dependencies]
+dev = [
+ { name = "apache-airflow-devel-common" },
+]
+mypy = [
+ { name = "apache-airflow-devel-common", extra = ["mypy"] },
+]
+
+[package.metadata]
+
+[package.metadata.requires-dev]
+dev = [{ name = "apache-airflow-devel-common", editable = "devel-common" }]
+mypy = [{ name = "apache-airflow-devel-common", extras = ["mypy"], editable =
"devel-common" }]
+
[[package]]
name = "apache-airflow-shared-template-rendering"
version = "0.0"
