This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 22538bf9593 [v3-4-test] Return 5xx on database deadlock when updating
task state (#75368) (#75652)
22538bf9593 is described below
commit 22538bf9593fd4ddd57fa7060756825b423457da
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Oct 26 19:26:00 2026 +0530
[v3-4-test] Return 5xx on database deadlock when updating task state
(#75368) (#75652)
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.
(cherry picked from commit 840d3a41c633abb1d5d97a8791af066b14b2a069)
Co-authored-by: Wei Lee <[email protected]>
---
.../execution_api/routes/task_instances.py | 7 +--
airflow-core/src/airflow/utils/sqlalchemy.py | 17 ++++++
.../versions/head/test_task_instances.py | 41 ++++++++++++++-
airflow-core/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 658c0e3055d..d69e51cde62 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
@@ -115,7 +115,7 @@ from airflow.models.xcom import build_xcom_read_query,
xcom_entity
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()
@@ -685,8 +685,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 18c5da08214..72b4ff02542 100644
--- a/airflow-core/src/airflow/utils/sqlalchemy.py
+++ b/airflow-core/src/airflow/utils/sqlalchemy.py
@@ -29,6 +29,7 @@ from uuid import UUID
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 8b26dd49b40..391d55aeed8 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 import StatusCode
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
@@ -3655,6 +3655,45 @@ class TestTIUpdateState:
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 5cb11e8cb41..6b1a9f3e64c 100644
--- a/airflow-core/tests/unit/utils/test_sqlalchemy.py
+++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py
@@ -27,7 +27,7 @@ import pytest
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 @@ from airflow.utils.sqlalchemy import (
compact_uuid_default,
ensure_pod_is_valid_after_unpickling,
get_dialect_name,
+ is_deadlock_error,
prohibit_commit,
with_row_locks,
)
@@ -523,3 +524,60 @@ class TestCompactUUID:
)
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
