jason810496 commented on code in PR #70983:
URL: https://github.com/apache/airflow/pull/70983#discussion_r3726632991
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py:
##########
@@ -41,16 +41,80 @@
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:
+ detail = f"Task instance not found for dag_id={dag_id!r},
run_id={dag_run_id!r}, task_id={task_id!r}"
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail=detail if map_index is None else f"{detail},
map_index={map_index}",
+ )
Review Comment:
We should show the `all_map_indices=True` in the error.
--
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]
