This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 87ac7048a0219e3c88e4d3db7b4eab91a8e71a10
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Fri Nov 6 16:44:26 2026 +0200
[v3-4-test] Return 404 when a DagRun disappears during task API writes
(#72907) (#76599)
* Return 404 instead of 500 when DagRun is missing in Execution API
In task_instances.py validate_inlets_and_outlets and task_state_store.py
set_task_state_store, missing DagRun instances triggered unhandled
AttributeError and ValueError exceptions that surfaced as HTTP 500
Internal Server Errors.
Catch the missing DagRun condition and return a clean 404 Not Found response
matching the pattern established in PR 72890.
* Distinguish missing runs from other state backend errors
Unrelated backend ValueErrors must not be reported as missing runs.
Exercise writes racing actual run deletion instead of manufacturing impossible
ORM results.
* Return 404 on task-state writes after a DagRun is deleted mid-request
* Fix task state store 404 regression coverage
The conflict merge kept the dedicated DagRunNotFound exception but retained
an upstream test that raised ValueError, so current-head CI failed before
exercising the route contract.
* Test the narrowed ValueError handling where it changed, in the public API
On main the Execution API route never caught ValueError, so the test there
passed without this change. The public API used to turn any ValueError into a
404.
---------
(cherry picked from commit f90d2607dd4714bff26d13951a5344a64e1057f1)
Signed-off-by: 1fanwang <[email protected]>
Signed-off-by: Stefan Wang <[email protected]>
Co-authored-by: Stefan Wang <[email protected]>
Co-authored-by: pierrejeambrun <[email protected]>
---
.../core_api/routes/public/task_state_store.py | 5 ++-
.../execution_api/routes/task_state_store.py | 17 +++++++-
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_10_30.py | 7 +++
airflow-core/src/airflow/state/metastore.py | 5 ++-
.../routes/public/test_task_state_store.py | 51 ++++++++++++++++++++--
.../versions/head/test_task_state_store.py | 38 +++++++++++++++-
.../versions/v2026_10_30/test_task_state_store.py | 51 ++++++++++++++++++++++
airflow-core/tests/unit/state/test_metastore.py | 5 ++-
9 files changed, 170 insertions(+), 11 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
index 0c58c87d3c0..f558700ca4b 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py
@@ -43,6 +43,7 @@ from
airflow.api_fastapi.core_api.services.public.task_coordinates import (
resolve_task_scope,
)
from airflow.configuration import conf
+from airflow.exceptions import DagRunNotFound
from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance as TI
@@ -284,7 +285,7 @@ def set_task_state_store(
expires_at = _resolve_expires_at(body.expires_at)
try:
_get_db_backend().set(scope, key, json.dumps(body.value),
expires_at=expires_at, session=session)
- except ValueError as e:
+ except DagRunNotFound as e:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND,
detail=str(e)) from e
@@ -324,7 +325,7 @@ def patch_task_state_store(
_get_db_backend().set(
scope, key, json.dumps(body.value),
expires_at=existing.expires_at, session=session
)
- except ValueError as e:
+ except DagRunNotFound as e:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND,
detail=str(e)) from e
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
index 94b4ec01391..749a8058a12 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
@@ -31,6 +31,8 @@ from
airflow.api_fastapi.execution_api.datamodels.task_state_store import (
TaskStateStoreResponse,
)
from airflow.api_fastapi.execution_api.security import ExecutionAPIRoute,
require_auth
+from airflow.api_fastapi.execution_api.versions.v2026_10_30 import
ReturnNotFoundForMissingDagRun
+from airflow.exceptions import DagRunNotFound
from airflow.models.taskinstance import TaskInstance as TI
from airflow.state import get_state_backend
@@ -93,7 +95,20 @@ def set_task_state_store(
) -> None:
"""Set a task state store key, creating or updating the row."""
scope = _get_task_scope_for_ti(task_instance_id, session)
- get_state_backend().set(scope, key, json.dumps(body.value),
expires_at=body.expires_at, session=session)
+ try:
+ get_state_backend().set(
+ scope, key, json.dumps(body.value), expires_at=body.expires_at,
session=session
+ )
+ except DagRunNotFound as e:
+ if not ReturnNotFoundForMissingDagRun.is_applied:
+ raise
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"DagRun with dag_id={scope.dag_id} and
run_id={scope.run_id} not found",
+ },
+ ) from e
@router.delete("/{task_instance_id}/{key:path}",
status_code=status.HTTP_204_NO_CONTENT)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index f67687e6e23..981b1279b52 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -65,6 +65,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
AddTaskInstanceRegionCoordinates,
AddTerminalStateRetryReasonField,
IdentifyArchivedTaskStateUpdates,
+ ReturnNotFoundForMissingDagRun,
)
from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import
(
AddPreviousIterationToXComFilterParams,
@@ -87,6 +88,7 @@ bundle = VersionBundle(
AddTaskInstanceRegionCoordinates,
IdentifyArchivedTaskStateUpdates,
AddFirstTaskRescheduleStartDateField,
+ ReturnNotFoundForMissingDagRun,
),
Version(
"2026-06-30",
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
index 65a0f0ef25b..2abb984dc73 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
@@ -85,6 +85,13 @@ class
IdentifyArchivedTaskStateUpdates(VersionChangeWithSideEffects):
instructions_to_migrate_to_previous_version = ()
+class ReturnNotFoundForMissingDagRun(VersionChangeWithSideEffects):
+ """Return 404 when a task state write cannot find its DagRun."""
+
+ description = __doc__
+ instructions_to_migrate_to_previous_version = ()
+
+
class AddArgBindingsToTIRunContext(VersionChangeWithSideEffects):
"""Add the ``arg_bindings`` argument-binding spec for stub
(foreign-runtime) tasks."""
diff --git a/airflow-core/src/airflow/state/metastore.py
b/airflow-core/src/airflow/state/metastore.py
index c3c367c3fe5..027901928ae 100644
--- a/airflow-core/src/airflow/state/metastore.py
+++ b/airflow-core/src/airflow/state/metastore.py
@@ -35,6 +35,7 @@ from airflow._shared.state import (
)
from airflow._shared.timezones import timezone
from airflow.configuration import conf
+from airflow.exceptions import DagRunNotFound
from airflow.models.asset_state_store import AssetStateStoreModel
from airflow.models.dagrun import DagRun
from airflow.models.task_state_store import TaskStateStoreModel
@@ -272,7 +273,7 @@ class MetastoreBackend(BaseStoreBackend):
)
)
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}")
+ raise DagRunNotFound(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,
@@ -499,7 +500,7 @@ class MetastoreBackend(BaseStoreBackend):
)
)
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}")
+ raise DagRunNotFound(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,
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
index bc6495b13e4..c224e66c1b0 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py
@@ -21,18 +21,21 @@ from unittest.mock import patch
import pytest
from pydantic import ValidationError
-from sqlalchemy import select
+from sqlalchemy import Select, delete, event, select
+from airflow import settings
from airflow._shared.timezones import timezone
from airflow.api_fastapi.core_api.datamodels.task_state_store import (
TaskStateStoreBody,
TaskStateStorePatchBody,
)
+from airflow.exceptions import DagRunNotFound
from airflow.models.dagrun import DagRun
from airflow.models.dynamic_region import DynamicRegion
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance
from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.utils.session import create_session
from airflow.utils.types import DagRunType
from tests_common.test_utils.config import conf_vars
@@ -222,6 +225,36 @@ class TestSetTaskState(TestTaskStateEndpoint):
response = test_client.put(bad_url, json={"value": "v"})
assert response.status_code == 404
+ @pytest.mark.parametrize("method", ["put", "patch"])
+ def test_write_dagrun_deleted_after_scope_lookup_returns_404(self,
test_client, method):
+ if method == "patch":
+ assert test_client.put(f"{BASE_URL}/job_id", json={"value":
"old"}).status_code == 204
+
+ dag_run_id = self.dag_run.id
+ deleted = False
+
+ def delete_run_before_lookup(conn, statement, multiparams, params,
execution_options):
+ nonlocal deleted
+ if (
+ deleted
+ or not isinstance(statement, Select)
+ or statement.column_descriptions[0]["expr"] is not DagRun.id
+ ):
+ return
+ deleted = True
+ with create_session(scoped=False) as cleanup_session:
+ cleanup_session.execute(delete(DagRun).where(DagRun.id ==
dag_run_id))
+
+ event.listen(settings.engine, "before_execute",
delete_run_before_lookup)
+ try:
+ response = getattr(test_client, method)(f"{BASE_URL}/job_id",
json={"value": "v"})
+ finally:
+ event.remove(settings.engine, "before_execute",
delete_run_before_lookup)
+
+ assert deleted
+ assert response.status_code == 404
+ assert response.json()["detail"] == f"No DagRun found for
dag_id={DAG_ID!r} run_id={RUN_ID!r}"
+
def test_set_nonexistent_task_id_returns_404(self, test_client):
"""set() returns 404 when task_id doesn not match any TaskInstance in
the run."""
bad_url =
f"/dags/{DAG_ID}/dagRuns/{RUN_ID}/taskInstances/nonexistent_task/state-store/job_id"
@@ -401,19 +434,31 @@ class TestPatchTaskState(TestTaskStateEndpoint):
assert row.value == expected_db
def test_patch_task_state_store_domain_error_returns_404(self,
test_client):
- """Domain-level ValueError raised during PATCH translates to HTTP
404."""
_create_task_state_store_row(self._session, "job_id", "initial",
self.dag_run)
self._session.commit()
with patch(
"airflow.state.metastore.MetastoreBackend.set",
autospec=True,
- side_effect=ValueError(f"No DagRun found for dag_id={DAG_ID!r}
run_id={RUN_ID!r}"),
+ side_effect=DagRunNotFound(f"No DagRun found for dag_id={DAG_ID!r}
run_id={RUN_ID!r}"),
):
response = test_client.patch(f"{BASE_URL}/job_id", json={"value":
"v2"})
assert response.status_code == 404
assert response.json()["detail"] == f"No DagRun found for
dag_id={DAG_ID!r} run_id={RUN_ID!r}"
+ @pytest.mark.parametrize("method", ["put", "patch"])
+ def test_unrelated_backend_value_error_is_not_not_found(self, test_client,
method):
+ _create_task_state_store_row(self._session, "job_id", "initial",
self.dag_run)
+ self._session.commit()
+
+ with patch(
+ "airflow.state.metastore.MetastoreBackend.set",
+ autospec=True,
+ side_effect=ValueError("Invalid backend configuration"),
+ ):
+ with pytest.raises(ValueError, match="Invalid backend
configuration"):
+ getattr(test_client, method)(f"{BASE_URL}/job_id",
json={"value": "v2"})
+
def test_unauthorized_returns_401(self, unauthenticated_test_client):
assert unauthenticated_test_client.patch(f"{BASE_URL}/job_id",
json={"value": "v"}).status_code == 401
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
index 042ae4ee653..52b1e306130 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
@@ -25,8 +25,9 @@ import pendulum
import pytest
from fastapi import Request
from fastapi.testclient import TestClient
-from sqlalchemy import delete, select
+from sqlalchemy import Select, delete, event, select
+from airflow import settings
from airflow._shared.timezones import timezone
from airflow.api_fastapi.app import cached_app
from airflow.api_fastapi.execution_api.datamodels.token import TIClaims,
TIToken
@@ -39,6 +40,8 @@ from airflow.utils.session import create_session
from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
from tests_common.pytest_plugin import CreateTaskInstance
@@ -253,6 +256,39 @@ class TestPutTaskState:
assert response.status_code == 404
+ def test_put_dagrun_deleted_after_scope_lookup_returns_404(
+ self, client: TestClient, session: Session, create_task_instance:
CreateTaskInstance
+ ):
+ ti = create_task_instance()
+ dag_run_id = ti.dag_run.id
+ session.commit()
+ deleted = False
+
+ def delete_run_before_lookup(conn, statement, multiparams, params,
execution_options):
+ nonlocal deleted
+ if (
+ deleted
+ or not isinstance(statement, Select)
+ or statement.column_descriptions[0]["expr"] is not DagRun.id
+ ):
+ return
+ deleted = True
+ with create_session(scoped=False) as cleanup_session:
+ cleanup_session.execute(delete(DagRun).where(DagRun.id ==
dag_run_id))
+
+ event.listen(settings.engine, "before_execute",
delete_run_before_lookup)
+ try:
+ response = client.put(_api_url(ti.id, "job_id"), json={"value":
"spark_001"})
+ finally:
+ event.remove(settings.engine, "before_execute",
delete_run_before_lookup)
+
+ assert deleted
+ assert response.status_code == 404
+ assert response.json()["detail"] == {
+ "reason": "not_found",
+ "message": f"DagRun with dag_id={ti.dag_id} and run_id={ti.run_id}
not found",
+ }
+
def test_put_key_with_slash(self, client: TestClient,
create_task_instance: CreateTaskInstance):
ti = create_task_instance()
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_state_store.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_state_store.py
new file mode 100644
index 00000000000..f71d935c05b
--- /dev/null
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_state_store.py
@@ -0,0 +1,51 @@
+# 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 unittest import mock
+
+import pytest
+
+from airflow.exceptions import DagRunNotFound
+
+pytestmark = pytest.mark.db_test
+
+
[email protected](
+ ("version", "expected_status"),
+ [
+ pytest.param("2026-06-30", 500, id="previous-version"),
+ pytest.param("2026-10-30", 404, id="new-version"),
+ ],
+)
[email protected]("airflow.api_fastapi.execution_api.routes.task_state_store.get_state_backend",
autospec=True)
+def test_missing_dagrun_response_by_version(
+ mock_get_backend, client, create_task_instance, monkeypatch, version,
expected_status
+):
+ ti = create_task_instance()
+ mock_get_backend.return_value.set.side_effect = DagRunNotFound("DagRun not
found")
+ client.headers["Airflow-API-Version"] = version
+ monkeypatch.setattr(client._transport, "raise_server_exceptions", False)
+
+ response = client.put(f"/execution/store/ti/{ti.id}/job_id",
json={"value": "spark_001"})
+
+ assert response.status_code == expected_status
+ if expected_status == 404:
+ assert response.json()["detail"] == {
+ "reason": "not_found",
+ "message": f"DagRun with dag_id={ti.dag_id} and run_id={ti.run_id}
not found",
+ }
diff --git a/airflow-core/tests/unit/state/test_metastore.py
b/airflow-core/tests/unit/state/test_metastore.py
index d353949374f..932bbf537a5 100644
--- a/airflow-core/tests/unit/state/test_metastore.py
+++ b/airflow-core/tests/unit/state/test_metastore.py
@@ -31,6 +31,7 @@ from airflow import settings
from airflow._shared.state import AssetStateStoreWriterKind
from airflow._shared.timezones import timezone
from airflow.configuration import conf
+from airflow.exceptions import DagRunNotFound
from airflow.models.asset import AssetModel
from airflow.models.asset_state_store import AssetStateStoreModel
from airflow.models.dagrun import DagRun, DagRunType
@@ -237,7 +238,7 @@ class TestMetastoreBackendTaskScope:
def test_set_raises_for_missing_dag_run(self, session: Session, backend:
MetastoreBackend):
scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run",
task_id=TASK_ID)
- with pytest.raises(ValueError, match="No DagRun found"):
+ with pytest.raises(DagRunNotFound, match="No DagRun found"):
backend.set(scope, "job_id", "app_1234", session=session)
def test_clear_scoped_to_map_index(self, session: Session, backend:
MetastoreBackend, dag_run: DagRun):
@@ -808,7 +809,7 @@ class TestMetastoreBackendAsync:
async def test_aset_task_raises_for_missing_dag_run(self, backend:
MetastoreBackend):
scope = TaskScope(dag_id="nonexistent_dag", run_id="nonexistent_run",
task_id=TASK_ID)
- with pytest.raises(ValueError, match="No DagRun found"):
+ with pytest.raises(DagRunNotFound, match="No DagRun found"):
await backend.aset(scope, "job_id", "app_async")
async def test_aset_and_aget_with_provided_session(
