This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun 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 f90d2607dd4 Return 404 when a DagRun disappears during task API writes
(#72907)
f90d2607dd4 is described below
commit f90d2607dd4714bff26d13951a5344a64e1057f1
Author: Stefan Wang <[email protected]>
AuthorDate: Fri Nov 6 04:20:46 2026 -0800
Return 404 when a DagRun disappears during task API writes (#72907)
* 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.
Signed-off-by: 1fanwang <[email protected]>
* 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.
Signed-off-by: 1fanwang <[email protected]>
* Return 404 on task-state writes after a DagRun is deleted mid-request
Signed-off-by: Stefan Wang <[email protected]>
* 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.
Signed-off-by: Stefan Wang <[email protected]>
* 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.
---------
Signed-off-by: 1fanwang <[email protected]>
Signed-off-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 4289af2be30..0b347e038ee 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
@@ -92,7 +94,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 89a82f696f5..2c169bbbe20 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
@@ -66,6 +66,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
AddTaskInstanceRegionCoordinates,
AddTerminalStateRetryReasonField,
IdentifyArchivedTaskStateUpdates,
+ ReturnNotFoundForMissingDagRun,
)
from airflow.api_fastapi.execution_api.versions.v2027_02_28 import (
AddDagProcessorJobEndpoints,
@@ -90,6 +91,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 6fef14d6bc9..7809e57bcf1 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
@@ -86,6 +86,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 f621c2b78fe..7c6e58b8877 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
@@ -275,6 +278,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 de9335dc83b..9b083b84a41 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):
@@ -815,7 +816,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(
