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