This is an automated email from the ASF dual-hosted git repository.
bbovenzi 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 f7e827b86e5 Return 404 from task state store endpoints for unknown
task instances (#70983)
f7e827b86e5 is described below
commit f7e827b86e51951304b6cdeccc1de7a04d14cce2
Author: Jyun-An Chen <[email protected]>
AuthorDate: Wed Aug 12 04:26:14 2026 +0800
Return 404 from task state store endpoints for unknown task instances
(#70983)
---
.../core_api/openapi/v2-rest-api-generated.yaml | 68 +++++-----
.../core_api/routes/public/task_state_store.py | 149 ++++++++++++---------
.../ui/openapi-gen/queries/ensureQueryData.ts | 2 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 2 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 10 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 2 +-
.../ui/openapi-gen/requests/services.gen.ts | 31 +++--
.../airflow/ui/openapi-gen/requests/types.gen.ts | 4 -
.../routes/public/test_task_state_store.py | 57 ++++++++
9 files changed, 198 insertions(+), 127 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 81b0a710c9e..b3ed8a3a2b8 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -6021,12 +6021,6 @@ paths:
schema:
$ref: '#/components/schemas/HTTPExceptionResponse'
description: Forbidden
- '404':
- content:
- application/json:
- schema:
- $ref: '#/components/schemas/HTTPExceptionResponse'
- description: Not Found
'422':
description: Validation Error
content:
@@ -6067,6 +6061,13 @@ paths:
schema:
type: string
title: Task Id
+ - name: all_map_indices
+ in: query
+ required: false
+ schema:
+ type: boolean
+ default: false
+ title: All Map Indices
- name: map_index
in: query
required: false
@@ -6075,13 +6076,6 @@ paths:
minimum: -1
default: -1
title: Map Index
- - name: all_map_indices
- in: query
- required: false
- schema:
- type: boolean
- default: false
- title: All Map Indices
responses:
'204':
description: Successful Response
@@ -6120,6 +6114,12 @@ paths:
- OAuth2PasswordBearer: []
- HTTPBearer: []
parameters:
+ - name: key
+ in: path
+ required: true
+ schema:
+ type: string
+ title: Key
- name: dag_id
in: path
required: true
@@ -6138,12 +6138,6 @@ paths:
schema:
type: string
title: Task Id
- - name: key
- in: path
- required: true
- schema:
- type: string
- title: Key
- name: map_index
in: query
required: false
@@ -6193,6 +6187,12 @@ paths:
- OAuth2PasswordBearer: []
- HTTPBearer: []
parameters:
+ - name: key
+ in: path
+ required: true
+ schema:
+ type: string
+ title: Key
- name: dag_id
in: path
required: true
@@ -6211,12 +6211,6 @@ paths:
schema:
type: string
title: Task Id
- - name: key
- in: path
- required: true
- schema:
- type: string
- title: Key
- name: map_index
in: query
required: false
@@ -6268,6 +6262,12 @@ paths:
- OAuth2PasswordBearer: []
- HTTPBearer: []
parameters:
+ - name: key
+ in: path
+ required: true
+ schema:
+ type: string
+ title: Key
- name: dag_id
in: path
required: true
@@ -6286,12 +6286,6 @@ paths:
schema:
type: string
title: Task Id
- - name: key
- in: path
- required: true
- schema:
- type: string
- title: Key
- name: map_index
in: query
required: false
@@ -6347,6 +6341,12 @@ paths:
- OAuth2PasswordBearer: []
- HTTPBearer: []
parameters:
+ - name: key
+ in: path
+ required: true
+ schema:
+ type: string
+ title: Key
- name: dag_id
in: path
required: true
@@ -6365,12 +6365,6 @@ paths:
schema:
type: string
title: Task Id
- - name: key
- in: path
- required: true
- schema:
- type: string
- title: Key
- name: map_index
in: query
required: false
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 54e1337215c..0d844b8b3ac 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
@@ -18,7 +18,7 @@ from __future__ import annotations
import json
from datetime import datetime, timedelta, timezone
-from typing import Annotated, Literal
+from typing import TYPE_CHECKING, Annotated, Literal
from fastapi import Depends, HTTPException, Query, status
from sqlalchemy import select
@@ -41,16 +41,83 @@ from airflow.models.task_state_store import
TaskStateStoreModel
from airflow.models.taskinstance import TaskInstance as TI
from airflow.state.metastore import _get_db_backend
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
task_state_store_router = AirflowRouter(
tags=["Task State Store"],
prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store",
)
-def _get_scope(dag_id: str, dag_run_id: str, task_id: str, map_index: int) ->
TaskScope:
+def _require_task_instance(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ map_index: int | None,
+ session: Session,
+) -> None:
+ """Raise 404 unless the task instance exists. ``map_index=None`` matches
any map index."""
+ statement = select(TI.task_id).where(
+ TI.dag_id == dag_id,
+ TI.run_id == dag_run_id,
+ TI.task_id == task_id,
+ )
+ if map_index is not None:
+ statement = statement.where(TI.map_index == map_index)
+ if session.scalar(statement.limit(1)) is None:
+ addressed_by = "all_map_indices=True" if map_index is None else
f"map_index={map_index}"
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail=(
+ f"Task instance not found for dag_id={dag_id!r},
run_id={dag_run_id!r}, "
+ f"task_id={task_id!r}, {addressed_by}"
+ ),
+ )
+
+
+def _resolve_scope(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ map_index: Annotated[int, Query(ge=-1)] = -1,
+) -> TaskScope:
+ """Map the path and query parameters onto the task instance they
address."""
return TaskScope(dag_id=dag_id, run_id=dag_run_id, task_id=task_id,
map_index=map_index)
+TaskScopeDep = Annotated[TaskScope, Depends(_resolve_scope)]
+
+
+def _validate_scope(scope: TaskScopeDep, session: SessionDep) -> TaskScope:
+ """Resolve the scope, 404ing when the task instance it addresses does not
exist."""
+ _require_task_instance(scope.dag_id, scope.run_id, scope.task_id,
scope.map_index, session)
+ return scope
+
+
+ValidatedTaskScopeDep = Annotated[TaskScope, Depends(_validate_scope)]
+
+
+def _validate_clear_scope(
+ scope: TaskScopeDep,
+ session: SessionDep,
+ all_map_indices: Annotated[bool, Query()] = False,
+) -> TaskScope:
+ """
+ Resolve the scope for a clear request, 404ing when the task instance does
not exist.
+
+ ``all_map_indices`` addresses the task across every index, so it is
validated against any
+ instance -- an expanded mapped task has no ``map_index=-1`` instance to
check.
+ """
+ _require_task_instance(
+ scope.dag_id, scope.run_id, scope.task_id, None if all_map_indices
else scope.map_index, session
+ )
+ return scope
+
+
+ValidatedClearTaskScopeDep = Annotated[TaskScope,
Depends(_validate_clear_scope)]
+
+
def _resolve_expires_at(expires_at: datetime | None | Literal["default"]) ->
datetime | None:
"""
Resolve the expires_at value from the request body.
@@ -72,35 +139,15 @@ def _resolve_expires_at(expires_at: datetime | None |
Literal["default"]) -> dat
return expires_at
-def _require_ti(dag_id: str, dag_run_id: str, task_id: str, map_index: int,
session: SessionDep) -> None:
- ti_exists = session.scalar(
- select(TI.task_id).where(
- TI.dag_id == dag_id,
- TI.run_id == dag_run_id,
- TI.task_id == task_id,
- TI.map_index == map_index,
- )
- )
- if ti_exists is None:
- raise HTTPException(
- status_code=status.HTTP_404_NOT_FOUND,
- detail=f"Task instance not found for dag_id={dag_id!r},
run_id={dag_run_id!r}, task_id={task_id!r}, map_index={map_index}",
- )
-
-
@task_state_store_router.get(
"",
- responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
dependencies=[Depends(requires_access_dag(method="GET",
access_entity=DagAccessEntity.TASK_INSTANCE))],
)
def list_task_state_store(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
+ scope: TaskScopeDep,
limit: QueryLimit,
offset: QueryOffset,
session: SessionDep,
- map_index: Annotated[int, Query(ge=-1)] = -1,
) -> TaskStateStoreCollectionResponse:
"""List all task state store entries for a task instance."""
base = (
@@ -111,10 +158,10 @@ def list_task_state_store(
TaskStateStoreModel.expires_at,
)
.where(
- TaskStateStoreModel.dag_id == dag_id,
- TaskStateStoreModel.run_id == dag_run_id,
- TaskStateStoreModel.task_id == task_id,
- TaskStateStoreModel.map_index == map_index,
+ TaskStateStoreModel.dag_id == scope.dag_id,
+ TaskStateStoreModel.run_id == scope.run_id,
+ TaskStateStoreModel.task_id == scope.task_id,
+ TaskStateStoreModel.map_index == scope.map_index,
)
.order_by(TaskStateStoreModel.key.asc())
)
@@ -142,12 +189,9 @@ def list_task_state_store(
dependencies=[Depends(requires_access_dag(method="GET",
access_entity=DagAccessEntity.TASK_INSTANCE))],
)
def get_task_state_store(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
+ scope: TaskScopeDep,
key: str,
session: SessionDep,
- map_index: Annotated[int, Query(ge=-1)] = -1,
) -> TaskStateStoreResponse:
"""Get a single task state store entry."""
row = session.execute(
@@ -157,10 +201,10 @@ def get_task_state_store(
TaskStateStoreModel.updated_at,
TaskStateStoreModel.expires_at,
).where(
- TaskStateStoreModel.dag_id == dag_id,
- TaskStateStoreModel.run_id == dag_run_id,
- TaskStateStoreModel.task_id == task_id,
- TaskStateStoreModel.map_index == map_index,
+ TaskStateStoreModel.dag_id == scope.dag_id,
+ TaskStateStoreModel.run_id == scope.run_id,
+ TaskStateStoreModel.task_id == scope.task_id,
+ TaskStateStoreModel.map_index == scope.map_index,
TaskStateStoreModel.key == key,
)
).one_or_none()
@@ -181,18 +225,13 @@ def get_task_state_store(
dependencies=[Depends(requires_access_dag(method="PUT",
access_entity=DagAccessEntity.TASK_INSTANCE))],
)
def set_task_state_store(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
+ scope: ValidatedTaskScopeDep,
key: str,
body: TaskStateStoreBody,
session: SessionDep,
- map_index: Annotated[int, Query(ge=-1)] = -1,
) -> None:
"""Set a task state store value. Creates or overwrites the key."""
- _require_ti(dag_id, dag_run_id, task_id, map_index, session)
expires_at = _resolve_expires_at(body.expires_at)
- scope = _get_scope(dag_id, dag_run_id, task_id, map_index)
try:
_get_db_backend().set(scope, key, json.dumps(body.value),
expires_at=expires_at, session=session)
except ValueError as e:
@@ -206,23 +245,18 @@ def set_task_state_store(
dependencies=[Depends(requires_access_dag(method="PUT",
access_entity=DagAccessEntity.TASK_INSTANCE))],
)
def patch_task_state_store(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
+ scope: ValidatedTaskScopeDep,
key: str,
body: TaskStateStorePatchBody,
session: SessionDep,
- map_index: Annotated[int, Query(ge=-1)] = -1,
) -> None:
"""Update the value of an existing task state store key."""
- _require_ti(dag_id, dag_run_id, task_id, map_index, session)
-
existing = session.execute(
select(TaskStateStoreModel.expires_at).where(
- TaskStateStoreModel.dag_id == dag_id,
- TaskStateStoreModel.run_id == dag_run_id,
- TaskStateStoreModel.task_id == task_id,
- TaskStateStoreModel.map_index == map_index,
+ TaskStateStoreModel.dag_id == scope.dag_id,
+ TaskStateStoreModel.run_id == scope.run_id,
+ TaskStateStoreModel.task_id == scope.task_id,
+ TaskStateStoreModel.map_index == scope.map_index,
TaskStateStoreModel.key == key,
)
).one_or_none()
@@ -233,7 +267,6 @@ def patch_task_state_store(
detail=f"Task state store key {key!r} not found",
)
- scope = _get_scope(dag_id, dag_run_id, task_id, map_index)
_get_db_backend().set(scope, key, json.dumps(body.value),
expires_at=existing.expires_at, session=session)
@@ -244,15 +277,11 @@ def patch_task_state_store(
dependencies=[Depends(requires_access_dag(method="DELETE",
access_entity=DagAccessEntity.TASK_INSTANCE))],
)
def delete_task_state_store(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
+ scope: ValidatedTaskScopeDep,
key: str,
session: SessionDep,
- map_index: Annotated[int, Query(ge=-1)] = -1,
) -> None:
"""Delete a single task state store key. No-op if the key does not
exist."""
- scope = _get_scope(dag_id, dag_run_id, task_id, map_index)
_get_db_backend().delete(scope, key, session=session)
@@ -263,11 +292,8 @@ def delete_task_state_store(
dependencies=[Depends(requires_access_dag(method="DELETE",
access_entity=DagAccessEntity.TASK_INSTANCE))],
)
def clear_task_state_store(
- dag_id: str,
- dag_run_id: str,
- task_id: str,
+ scope: ValidatedClearTaskScopeDep,
session: SessionDep,
- map_index: Annotated[int, Query(ge=-1)] = -1,
all_map_indices: Annotated[bool, Query()] = False,
) -> None:
"""
@@ -276,5 +302,4 @@ def clear_task_state_store(
When ``all_map_indices=true``, state store is cleared for every map index
of the task and
the ``map_index`` parameter is ignored.
"""
- scope = _get_scope(dag_id, dag_run_id, task_id, map_index)
_get_db_backend().clear(scope, all_map_indices=all_map_indices,
session=session)
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index 637e37a62f4..c7d379c85d9 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1583,10 +1583,10 @@ export const
ensureUseTaskStateStoreServiceListTaskStateStoreData = (queryClient
* Get Task State Store
* Get a single task state store entry.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.mapIndex
* @returns TaskStateStoreResponse Successful Response
* @throws ApiError
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index d7d3fecb52c..f0af02f94f6 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1583,10 +1583,10 @@ export const
prefetchUseTaskStateStoreServiceListTaskStateStore = (queryClient:
* Get Task State Store
* Get a single task state store entry.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.mapIndex
* @returns TaskStateStoreResponse Successful Response
* @throws ApiError
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index 1b993afa9b0..0c69edfd97b 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1583,10 +1583,10 @@ export const useTaskStateStoreServiceListTaskStateStore
= <TData = Common.TaskSt
* Get Task State Store
* Get a single task state store entry.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.mapIndex
* @returns TaskStateStoreResponse Successful Response
* @throws ApiError
@@ -2509,10 +2509,10 @@ export const
useAssetStateStoreServiceSetAssetStateStore = <TData = Common.Asset
* Set Task State Store
* Set a task state store value. Creates or overwrites the key.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.requestBody
* @param data.mapIndex
* @returns void Successful Response
@@ -2921,10 +2921,10 @@ export const usePoolServiceBulkPools = <TData =
Common.PoolServiceBulkPoolsMutat
* Patch Task State Store
* Update the value of an existing task state store key.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.requestBody
* @param data.mapIndex
* @returns unknown Successful Response
@@ -3168,8 +3168,8 @@ export const
useAssetStateStoreServiceDeleteAssetStateStore = <TData = Common.As
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.mapIndex
* @param data.allMapIndices
+* @param data.mapIndex
* @returns void Successful Response
* @throws ApiError
*/
@@ -3190,10 +3190,10 @@ export const
useTaskStateStoreServiceClearTaskStateStore = <TData = Common.TaskS
* Delete Task State Store
* Delete a single task state store key. No-op if the key does not exist.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.mapIndex
* @returns void Successful Response
* @throws ApiError
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index 4d2dc42adca..ec1e3504f7d 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1583,10 +1583,10 @@ export const
useTaskStateStoreServiceListTaskStateStoreSuspense = <TData = Commo
* Get Task State Store
* Get a single task state store entry.
* @param data The data for the request.
+* @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
-* @param data.key
* @param data.mapIndex
* @returns TaskStateStoreResponse Successful Response
* @throws ApiError
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 09e450335b1..0fc5c1c220f 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -3891,7 +3891,6 @@ export class TaskStateStoreService {
errors: {
401: 'Unauthorized',
403: 'Forbidden',
- 404: 'Not Found',
422: 'Validation Error'
}
});
@@ -3907,8 +3906,8 @@ export class TaskStateStoreService {
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
- * @param data.mapIndex
* @param data.allMapIndices
+ * @param data.mapIndex
* @returns void Successful Response
* @throws ApiError
*/
@@ -3922,8 +3921,8 @@ export class TaskStateStoreService {
task_id: data.taskId
},
query: {
- map_index: data.mapIndex,
- all_map_indices: data.allMapIndices
+ all_map_indices: data.allMapIndices,
+ map_index: data.mapIndex
},
errors: {
401: 'Unauthorized',
@@ -3938,10 +3937,10 @@ export class TaskStateStoreService {
* Get Task State Store
* Get a single task state store entry.
* @param data The data for the request.
+ * @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
- * @param data.key
* @param data.mapIndex
* @returns TaskStateStoreResponse Successful Response
* @throws ApiError
@@ -3951,10 +3950,10 @@ export class TaskStateStoreService {
method: 'GET',
url:
'/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
path: {
+ key: data.key,
dag_id: data.dagId,
dag_run_id: data.dagRunId,
- task_id: data.taskId,
- key: data.key
+ task_id: data.taskId
},
query: {
map_index: data.mapIndex
@@ -3972,10 +3971,10 @@ export class TaskStateStoreService {
* Set Task State Store
* Set a task state store value. Creates or overwrites the key.
* @param data The data for the request.
+ * @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
- * @param data.key
* @param data.requestBody
* @param data.mapIndex
* @returns void Successful Response
@@ -3986,10 +3985,10 @@ export class TaskStateStoreService {
method: 'PUT',
url:
'/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
path: {
+ key: data.key,
dag_id: data.dagId,
dag_run_id: data.dagRunId,
- task_id: data.taskId,
- key: data.key
+ task_id: data.taskId
},
query: {
map_index: data.mapIndex
@@ -4009,10 +4008,10 @@ export class TaskStateStoreService {
* Patch Task State Store
* Update the value of an existing task state store key.
* @param data The data for the request.
+ * @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
- * @param data.key
* @param data.requestBody
* @param data.mapIndex
* @returns unknown Successful Response
@@ -4023,10 +4022,10 @@ export class TaskStateStoreService {
method: 'PATCH',
url:
'/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
path: {
+ key: data.key,
dag_id: data.dagId,
dag_run_id: data.dagRunId,
- task_id: data.taskId,
- key: data.key
+ task_id: data.taskId
},
query: {
map_index: data.mapIndex
@@ -4046,10 +4045,10 @@ export class TaskStateStoreService {
* Delete Task State Store
* Delete a single task state store key. No-op if the key does not exist.
* @param data The data for the request.
+ * @param data.key
* @param data.dagId
* @param data.dagRunId
* @param data.taskId
- * @param data.key
* @param data.mapIndex
* @returns void Successful Response
* @throws ApiError
@@ -4059,10 +4058,10 @@ export class TaskStateStoreService {
method: 'DELETE',
url:
'/api/v2/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/state-store/{key}',
path: {
+ key: data.key,
dag_id: data.dagId,
dag_run_id: data.dagRunId,
- task_id: data.taskId,
- key: data.key
+ task_id: data.taskId
},
query: {
map_index: data.mapIndex
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index 13252e1126f..92b30f77199 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -7803,10 +7803,6 @@ export type $OpenApiTs = {
* Forbidden
*/
403: HTTPExceptionResponse;
- /**
- * Not Found
- */
- 404: HTTPExceptionResponse;
/**
* Validation Error
*/
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 650858656da..ac117b31d04 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
@@ -498,3 +498,60 @@ class
TestRoutesNeverCallCustomBackend(TestTaskStateEndpoint):
getattr(test_client, method)(path, **kwargs)
mock_get_backend.assert_not_called()
+
+
+class TestUnknownTaskInstance(TestTaskStateEndpoint):
+ UNKNOWN_TASK_URL =
f"/dags/{DAG_ID}/dagRuns/{RUN_ID}/taskInstances/no_such_task/state-store"
+ UNKNOWN_RUN_URL =
f"/dags/{DAG_ID}/dagRuns/no_such_run/taskInstances/{TASK_ID}/state-store"
+
+ @pytest.mark.parametrize(
+ ("method", "suffix", "kwargs"),
+ [
+ ("put", "/job_id", {"json": {"value": "v"}}),
+ ("patch", "/job_id", {"json": {"value": "v"}}),
+ ("delete", "/job_id", {}),
+ ("delete", "", {}),
+ ("delete", "?all_map_indices=true", {}),
+ ],
+ )
+ @pytest.mark.parametrize("base", [UNKNOWN_TASK_URL, UNKNOWN_RUN_URL],
ids=["task_id", "run_id"])
+ def test_returns_404(self, test_client, base, method, suffix, kwargs):
+ response = getattr(test_client, method)(f"{base}{suffix}", **kwargs)
+ assert response.status_code == 404
+ assert "Task instance not found" in response.json()["detail"]
+
+ @pytest.mark.parametrize(
+ ("method", "suffix", "kwargs"),
+ [
+ ("put", "/job_id", {"json": {"value": "v"}}),
+ ("patch", "/job_id", {"json": {"value": "v"}}),
+ ("delete", "/job_id", {}),
+ ("delete", "", {}),
+ ],
+ )
+ def test_unknown_map_index_returns_404(self, test_client, method, suffix,
kwargs):
+ response = getattr(test_client,
method)(f"{BASE_URL}{suffix}?map_index=7", **kwargs)
+ assert response.status_code == 404
+ assert "map_index=7" in response.json()["detail"]
+
+ def test_clear_all_map_indices_reports_the_indices_it_looked_for(self,
test_client):
+ response =
test_client.delete(f"{self.UNKNOWN_TASK_URL}?all_map_indices=true")
+ assert response.status_code == 404
+ assert response.json()["detail"].endswith("task_id='no_such_task',
all_map_indices=True")
+
+ @pytest.mark.parametrize("base", [UNKNOWN_TASK_URL, UNKNOWN_RUN_URL],
ids=["task_id", "run_id"])
+ def test_list_returns_an_empty_collection(self, test_client, base):
+ response = test_client.get(base)
+ assert response.status_code == 200
+ assert response.json() == {"task_state_store": [], "total_entries": 0}
+
+ @pytest.mark.parametrize("base", [UNKNOWN_TASK_URL, UNKNOWN_RUN_URL],
ids=["task_id", "run_id"])
+ def test_get_reports_the_missing_key(self, test_client, base):
+ response = test_client.get(f"{base}/job_id")
+ assert response.status_code == 404
+ assert response.json()["detail"] == "Task state store key 'job_id' not
found"
+
+ def test_missing_key_still_reports_the_key_not_the_task_instance(self,
test_client):
+ response = test_client.get(f"{BASE_URL}/no_such_key")
+ assert response.status_code == 404
+ assert response.json()["detail"] == "Task state store key
'no_such_key' not found"
