pierrejeambrun commented on code in PR #72907:
URL: https://github.com/apache/airflow/pull/72907#discussion_r4491784753
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py:
##########
@@ -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=scope, key=key, value=json.dumps(body.value),
expires_at=body.expires_at, session=session
Review Comment:
nit: could we keep the original positional call here? Only the
`try`/`except` is needed for the fix, and it keeps the diff smaller.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py:
##########
@@ -275,6 +279,49 @@ def test_put_missing_ti_returns_404(self, client:
TestClient):
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",
+ }
+
+
@mock.patch("airflow.api_fastapi.execution_api.routes.task_state_store.get_state_backend",
autospec=True)
+ def test_put_unrelated_backend_value_error_is_not_not_found(
Review Comment:
I think this one passes on main too, since the Execution API route never
caught `ValueError`. The place where unrelated `ValueError`s changed behavior
is the public API (they used to become 404). Could we move this test to
`core_api/routes/public/test_task_state_store.py` instead, so it covers the
narrowed `except`?
--
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]
