From 5dd46651da6cd7fcda15452abd6bc9ac45c27c24 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Thu, 22 Oct 2026 10:39:52 +0800 Subject: [PATCH] Return 5xx on database deadlock when updating task state A deadlock is transient and the task would succeed on a retry, but the catch-all around the state update marked the task instance FAILED, which is permanent. The Task SDK client already retries 5xx responses from the Execution API, so letting the deadlock surface as a 500 gives the worker that retry. Only deadlocks are treated this way. Recognising other lock errors would change how existing callers of is_lock_not_available_error behave. --- .../execution_api/routes/task_instances.py | 7 ++- airflow-core/src/airflow/utils/sqlalchemy.py | 17 ++++++ .../versions/head/test_task_instances.py | 41 ++++++++++++- .../tests/unit/utils/test_sqlalchemy.py | 60 ++++++++++++++++++- 4 files changed, 120 insertions(+), 5 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index d092b9f62a63e..4059a3ba4ce72 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -120,7 +120,7 @@ from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent -from airflow.utils.sqlalchemy import get_dialect_name +from airflow.utils.sqlalchemy import get_dialect_name, is_deadlock_error from airflow.utils.state import DagRunState, IntermediateTIState, TaskInstanceState, TerminalTIState router = VersionedAPIRouter(route_class=ExecutionAPIRoute) @@ -693,8 +693,9 @@ def ti_update_state( except DataError: # Let DataErrorHandler return a 422 instead of silently marking the TI FAILED below. raise - except Exception: - if loop_gate is not None: + except Exception as e: + # A deadlock is transient: return a 5xx so the Task SDK client retries instead of failing the TI. + if loop_gate is not None or is_deadlock_error(e): raise # Set a task to failed in case any unexpected exception happened during task state update log.exception( diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index 18c5da082145f..72b4ff02542b6 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -29,6 +29,7 @@ from sqlalchemy import TIMESTAMP, PickleType, String, Uuid, event, nullsfirst, text from sqlalchemy.dialects import mysql from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.exc import DBAPIError from sqlalchemy.ext.compiler import compiles from sqlalchemy.sql.expression import ColumnElement from sqlalchemy.sql.functions import FunctionElement @@ -772,6 +773,22 @@ def is_lock_not_available_error(error: OperationalError): return False +def is_deadlock_error(error: Exception) -> bool: + """Check if the Error is a database deadlock.""" + if not isinstance(error, DBAPIError) or error.orig is None: + return False + # DB specific error codes: + # Postgres: 40P01 (deadlock_detected) + # MySQL: 1213, 'Deadlock found when trying to get lock; try restarting transaction' + # psycopg2 exposes `pgcode`, psycopg3 exposes `sqlstate` + db_err_code = ( + getattr(error.orig, "pgcode", None) + or getattr(error.orig, "sqlstate", None) + or (error.orig.args[0] if error.orig.args else None) + ) + return db_err_code in ("40P01", 1213) + + def make_dialect_kwarg(dialect: str) -> dict[str, str | Iterable[str]]: """Create an SQLAlchemy-version-aware dialect keyword argument.""" return {"dialect_names": (dialect,)} diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 1b320c462c439..aef15aade0b3f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -38,7 +38,7 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from pydantic import ValidationError from sqlalchemy import delete, select, update -from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.exc import OperationalError, SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session from sqlalchemy.orm.exc import StaleDataError @@ -3730,6 +3730,45 @@ def fail_ti_update(db_session, statement, *args, **kwargs): assert isinstance(detail, dict) assert detail.get("reason") == "Database error" + @pytest.mark.parametrize( + ("pgcode", "expected_status", "expected_state"), + [ + pytest.param("40P01", 500, State.RUNNING, id="deadlock-is-retryable"), + pytest.param("55P03", 204, State.FAILED, id="other-operational-error-fails-ti"), + pytest.param("psycopg3", 500, State.RUNNING, id="psycopg3-deadlock-is-retryable"), + ], + ) + def test_ti_update_state_deadlock_returns_5xx_instead_of_failing( + self, client, session, create_task_instance, mocker, pgcode, expected_status, expected_state + ): + ti = create_task_instance( + task_id="test_ti_update_state_deadlock_returns_5xx_instead_of_failing", + state=State.RUNNING, + ) + session.commit() + if pgcode == "psycopg3": + errors = pytest.importorskip( + "psycopg.errors", reason="Test only runs when psycopg v3 is installed." + ) + orig = errors.DeadlockDetected("deadlock detected") + else: + orig = mock.Mock(sqlstate=None, pgcode=pgcode, args=("db error",)) + mocker.patch.object( + task_instances_route, + "_create_ti_state_update_query_and_update_state", + autospec=True, + side_effect=OperationalError("statement", {}, orig), + ) + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()}, + ) + + assert response.status_code == expected_status + session.expire_all() + assert session.get(TaskInstance, ti.id).state == expected_state + def test_ti_run_database_error(self, client, session, create_task_instance): """ Test that a database error is handled correctly when starting the Task Instance. diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 5cb11e8cb4143..6b1a9f3e64c28 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -27,7 +27,7 @@ from kubernetes.client import Configuration, models as k8s from sqlalchemy import Column, MetaData, Table, text from sqlalchemy.dialects import mysql, postgresql, sqlite -from sqlalchemy.exc import StatementError +from sqlalchemy.exc import OperationalError, StatementError from sqlalchemy.schema import CreateTable from airflow import settings @@ -43,6 +43,7 @@ compact_uuid_default, ensure_pod_is_valid_after_unpickling, get_dialect_name, + is_deadlock_error, prohibit_commit, with_row_locks, ) @@ -523,3 +524,60 @@ def test_mysql_column_default_is_parenthesized_without_a_server_version(self): ) ddl = str(CreateTable(table).compile(dialect=mysql.dialect())) assert "DEFAULT (UNHEX('0190a1b2c3d47e5f8a9b0c1d2e3f4a5b'))" in ddl + + +class TestIsDeadlockError: + @staticmethod + def _make_error(orig): + return OperationalError("statement", {}, orig) + + @pytest.mark.parametrize( + ("orig", "expected"), + [ + pytest.param( + mock.Mock(sqlstate=None, pgcode="40P01", args=("deadlock detected",)), + True, + id="postgres-deadlock", + ), + pytest.param( + mock.Mock(sqlstate=None, pgcode=None, args=(1213, "Deadlock found when trying to get lock")), + True, + id="mysql-deadlock", + ), + pytest.param( + mock.Mock(sqlstate=None, pgcode="55P03", args=("lock not available",)), + False, + id="postgres-nowait", + ), + pytest.param( + mock.Mock(sqlstate=None, pgcode=None, args=(1205, "Lock wait timeout exceeded")), + False, + id="mysql-lock-wait-timeout", + ), + pytest.param( + mock.Mock(sqlstate=None, pgcode=None, args=("connection reset",)), + False, + id="other-operational", + ), + ], + ) + def test_is_deadlock_error(self, orig, expected): + assert is_deadlock_error(self._make_error(orig)) is expected + + @pytest.mark.parametrize( + ("error_name", "expected"), + [ + pytest.param("DeadlockDetected", True, id="psycopg3-deadlock"), + pytest.param("LockNotAvailable", False, id="psycopg3-nowait"), + ], + ) + def test_is_deadlock_error_psycopg3(self, error_name, expected): + errors = pytest.importorskip("psycopg.errors", reason="Test only runs when psycopg v3 is installed.") + orig = getattr(errors, error_name)("db error") + assert is_deadlock_error(self._make_error(orig)) is expected + + def test_dbapi_error_without_orig_is_not_deadlock(self): + assert is_deadlock_error(OperationalError("statement", {}, None)) is False + + def test_non_dbapi_error_is_not_deadlock(self): + assert is_deadlock_error(RuntimeError("deadlock")) is False