ashb commented on code in PR #65759: URL: https://github.com/apache/airflow/pull/65759#discussion_r3138882327
########## airflow-core/tests/unit/models/test_task_state.py: ########## @@ -0,0 +1,234 @@ +# +# 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", + updated_by_try=1, + ) + 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.updated_by_try == 1 + 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): + """Validates that retries share state.""" + 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", + updated_by_try=1, + ) + session.add(row) + session.flush() + + row.value = "application_try1" Review Comment: `row.value = "application_try1"` is the same value as the initial write on line 132. The test only asserts `updated_by_try` changed but doesn't verify the value was actually updated by the retry. Should be something like `row.value = "application_try2"` with a corresponding `assert result.value == "application_try2"`. ########## shared/state/src/airflow_shared/state/__init__.py: ########## @@ -0,0 +1,79 @@ +# 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).""" + + 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: str + + +StateScope = TaskScope | AssetScope + + +class BaseStateBackend(ABC): + """ + Abstract backend for reading and writing task and asset state. + + - Sync methods are abstract. + - Async methods default to calling their sync counterparts so that sync-only backends (e.g. the default metastore backend) + get async support for free. Backends with native async I/O (e.g. S3) should + override the async methods. + """ + + @abstractmethod + def get(self, scope: StateScope, key: str) -> str | None: + """Return the stored value, or None if the key does not exist.""" + + @abstractmethod + def set(self, scope: StateScope, key: str, value: str) -> None: + """Write or overwrite the value for the given key.""" + + @abstractmethod + def delete(self, scope: StateScope, key: str) -> None: + """Delete a single key. No-op if the key does not exist.""" + + @abstractmethod + def clear(self, scope: StateScope) -> None: + """Delete all keys under the given scope.""" + + async def aget(self, scope: StateScope, key: str) -> str | None: + return self.get(scope, key) + + async def aset(self, scope: StateScope, key: str, value: str) -> None: + self.set(scope, key, value) + + async def adelete(self, scope: StateScope, key: str) -> None: + self.delete(scope, key) + + async def aclear(self, scope: StateScope) -> None: + self.clear(scope) Review Comment: The async impls call the sync methods directly, which blocks the event loop. This will be a problem for triggers, which run in an async event loop and would call the backend directly (not via the Execution API). These need native async implementations — at minimum the MetastoreStateBackend needs async SQLAlchemy session support for the async variants. The base class should probably make the async methods abstract too (alongside the sync ones), so backend authors are forced to provide real async implementations rather than silently inheriting blocking defaults. The PR description says these are "required for AIP-98 compatibility — async tasks must be able to call `await task_state.aset(...)` without blocking the event loop" — but as written they do block. ########## shared/state/src/airflow_shared/state/__init__.py: ########## @@ -0,0 +1,79 @@ +# 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).""" + + 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.""" + Review Comment: `asset_id` is an integer FK in the DB but typed as `str` here. This forces `int(scope.asset_id)` casts in every asset method in MetastoreStateBackend. If the field is always numeric, typing it as `int` removes the scattered casts. If it needs to be `str` for serialization across the Execution API (PR 2), a comment explaining why would help. ########## airflow-core/src/airflow/config_templates/config.yml: ########## @@ -2987,7 +2987,19 @@ dag_processor: type: string example: ~ default: "warning" - +state: Review Comment: nit: missing blank line before `state:` — other top-level sections in this file are separated by blank lines. Same between `state:` and `profiling:` below. ########## 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.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) + 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(String(ID_LEN, **COLLATION_ARGS), nullable=False) + run_id: Mapped[str] = mapped_column(String(ID_LEN, **COLLATION_ARGS), nullable=False) + + value: Mapped[str] = mapped_column(Text, nullable=False) + updated_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False) + updated_by_try: Mapped[int | None] = mapped_column(Integer, nullable=True) Review Comment: `updated_by_try` is never set by `MetastoreStateBackend._set_task`, and `TaskScope` doesn't carry a try number — so this column will always be NULL for rows written through the backend. More broadly: should this be the `TaskInstance.id` UUID instead of an integer try number? The worker's JWT is scoped to `task_instance_id` (a UUID), so that's what the caller naturally has. A UUID is globally unique and directly joinable to `task_instance` / `task_instance_history`, whereas an integer `try_number` is only meaningful combined with the rest of the PK scope you already have. Something like `updated_by_ti_id` (UUID FK to `task_instance.id`) would be more useful for auditing and wouldn't require an extra lookup on the write path. ########## 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-statse Review Comment: Typo: ```suggestion name: Run mypy for shared-state ``` ########## airflow-core/src/airflow/state/metastore.py: ########## @@ -0,0 +1,162 @@ +# +# 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, provide_session + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + +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: + if isinstance(scope, TaskScope): + return self._get_task(scope, key, session=session) + return self._get_asset(scope, key, session=session) + + @provide_session + def set(self, scope: StateScope, key: str, value: str, *, session: Session = NEW_SESSION) -> None: + if isinstance(scope, TaskScope): + self._set_task(scope, key, value, session=session) + else: + self._set_asset(scope, key, value, session=session) + + @provide_session + def delete(self, scope: StateScope, key: str, *, session: Session = NEW_SESSION) -> None: + if isinstance(scope, TaskScope): + self._delete_task(scope, key, session=session) + else: + self._delete_asset(scope, key, session=session) + + @provide_session + def clear(self, scope: StateScope, *, session: Session = NEW_SESSION) -> None: + if isinstance(scope, TaskScope): + self._clear_task(scope, session=session) + else: + self._clear_asset(scope, session=session) + + def _get_task(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(self, scope: TaskScope, key: str, value: str, *, session: Session) -> 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, + ) + ) + if row is None: + dag_run_id = session.scalar( + select(DagRun.id).where( + DagRun.dag_id == scope.dag_id, + DagRun.run_id == scope.run_id, + ) + ) + row = TaskStateModel( + 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, + ) + session.add(row) + row.value = value + row.updated_at = timezone.utcnow() Review Comment: SELECT-then-INSERT without locking. Two concurrent calls for the same `(scope, key)` can both see `None`, both INSERT, and one gets an IntegrityError. Same pattern in `_set_asset`. A database-level UPSERT (`INSERT...ON CONFLICT UPDATE`) would be atomic and also more efficient (single round-trip instead of two). In practice concurrent writes to the same key should be rare, but this is easy to get right from the start. -- 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]
