Return 5xx on database deadlock when updating task state

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(

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,)}

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.

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