Do not send SKIP LOCKED to servers that do not implement it

airflow-core/newsfragments/71250.bugfix.rst

@@ -0,0 +1 @@
+Airflow no longer sends ``SKIP LOCKED`` to MySQL-compatible servers that accept the clause and then ignore it. Such a server hands the same rows to every scheduler at once, silently breaking the exclusivity that the scheduler relies on when claiming task instances. Those servers now get a plain blocking ``FOR UPDATE`` instead, and a warning is logged once. MySQL and PostgreSQL are unaffected.

airflow-core/src/airflow/utils/sqlalchemy.py

@@ -24,6 +24,7 @@
 import logging
 from collections.abc import Generator
 from typing import TYPE_CHECKING, Any
+from weakref import WeakKeyDictionary
 
 from sqlalchemy import TIMESTAMP, PickleType, String, event, nullsfirst, text
 from sqlalchemy.dialects import mysql
@@ -38,12 +39,13 @@
 from airflow.serialization.enums import Encoding
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, MutableMapping
 
     from kubernetes.client.models.v1_pod import V1Pod
     from sqlalchemy.dialects.mysql.dml import Insert as MySQLInsert
     from sqlalchemy.dialects.postgresql.dml import Insert as PostgreSQLInsert
     from sqlalchemy.dialects.sqlite.dml import Insert as SQLiteInsert
+    from sqlalchemy.engine import Engine
     from sqlalchemy.exc import OperationalError
     from sqlalchemy.orm import Session
     from sqlalchemy.sql import Select
@@ -518,6 +520,43 @@ def nulls_first(col: ColumnElement, session: Session) -> ColumnElement:
 
 USE_ROW_LEVEL_LOCKING: bool = conf.getboolean("scheduler", "use_row_level_locking", fallback=True)
 
+_SKIP_LOCKED_IGNORED_BY = ("tidb",)
+
+_skip_locked_support: MutableMapping[Engine, bool] = WeakKeyDictionary()
+
+
+def _honors_skip_locked(engine: Engine) -> bool:
+    """
+    Whether the server actually skips locked rows when asked to.
+
+    TiDB parses ``SKIP LOCKED``, silently drops it, and emits no warning
+    (https://github.com/pingcap/tidb/issues/18207). Callers use ``SKIP LOCKED``
+    to claim work exclusively, so a server that ignores it hands the same rows
+    to every scheduler at once instead of failing.
+    """
+    cached = _skip_locked_support.get(engine)
+    if cached is not None:
+        return cached
+
+    honored = True
+    if engine.dialect.name == "mysql":
+        try:
+            with engine.connect() as conn:
+                banner = str(conn.exec_driver_sql("SELECT VERSION()").scalar() or "").lower()
+        except Exception:
+            log.debug("Could not read the server version banner; assuming SKIP LOCKED works", exc_info=True)
+            banner = ""
+        if any(name in banner for name in _SKIP_LOCKED_IGNORED_BY):
+            honored = False
+            log.warning(
+                "Database server reports as %r, which accepts SKIP LOCKED but does not honor it. "
+                "Falling back to plain FOR UPDATE so concurrent schedulers cannot claim the same "
+                "rows. Schedulers will block on each other instead of skipping ahead.",
+                banner,
+            )
+    _skip_locked_support[engine] = honored
+    return honored
+
 
 def with_row_locks(
     query: Select,
@@ -557,13 +596,14 @@ def with_row_locks(
     # Don't use row level locks if the MySQL dialect (Mariadb & MySQL < 8) does not support it.
     if not USE_ROW_LEVEL_LOCKING:
         return query
+    bind = session.bind
     if dialect_name == "mysql" and not getattr(
-        session.bind.dialect if session.bind else None, "supports_for_update_of", False
+        bind.dialect if bind else None, "supports_for_update_of", False
     ):
         return query
     if nowait:
         kwargs["nowait"] = True
-    if skip_locked:
+    if skip_locked and (bind is None or _honors_skip_locked(bind.engine)):
         kwargs["skip_locked"] = True
     if key_share:
         kwargs["key_share"] = True

airflow-core/tests/unit/utils/test_sqlalchemy.py

@@ -180,6 +180,89 @@ def test_with_row_locks(
             assert returned_value == query
             query.with_for_update.assert_not_called()
 
+    @staticmethod
+    def _mysql_bind(banner: str):
+        """Build a bind whose server reports ``banner`` from ``SELECT VERSION()``."""
+
+        class _Dialect:
+            name = "mysql"
+            supports_for_update_of = True
+
+        conn = mock.MagicMock()
+        conn.exec_driver_sql.return_value.scalar.return_value = banner
+        engine = mock.MagicMock()
+        engine.dialect = _Dialect()
+        engine.connect.return_value.__enter__.return_value = conn
+        bind = mock.MagicMock()
+        bind.dialect = engine.dialect
+        bind.engine = engine
+        return bind, engine
+
+    @pytest.mark.parametrize(
+        ("banner", "expected_skip_locked"),
+        [
+            pytest.param("8.0.11-TiDB-v8.5.1", False, id="tidb-ignores-skip-locked"),
+            pytest.param("8.0.11-tidb-v7.5.0", False, id="tidb-lowercase-banner"),
+            pytest.param("8.4.0", True, id="mysql-84"),
+            pytest.param("8.0.39-0ubuntu0.22.04.1", True, id="mysql-80"),
+        ],
+    )
+    def test_with_row_locks_skip_locked_only_when_server_honors_it(self, banner, expected_skip_locked):
+        """TiDB accepts SKIP LOCKED and discards it, so we must not send it there."""
+        query = mock.Mock()
+        bind, _ = self._mysql_bind(banner)
+        session = mock.Mock()
+        session.bind = bind
+        session.get_bind.return_value = bind
+
+        with mock.patch("airflow.utils.sqlalchemy.USE_ROW_LEVEL_LOCKING", True):
+            with_row_locks(query=query, session=session, skip_locked=True)
+
+        expected = {"key_share": True}
+        if expected_skip_locked:
+            expected["skip_locked"] = True
+        query.with_for_update.assert_called_once_with(**expected)
+
+    def test_with_row_locks_probes_the_server_version_once(self):
+        query = mock.Mock()
+        bind, engine = self._mysql_bind("8.0.11-TiDB-v8.5.1")
+        session = mock.Mock()
+        session.bind = bind
+        session.get_bind.return_value = bind
+
+        with mock.patch("airflow.utils.sqlalchemy.USE_ROW_LEVEL_LOCKING", True):
+            for _ in range(3):
+                with_row_locks(query=query, session=session, skip_locked=True)
+
+        assert engine.connect.call_count == 1
+
+    def test_with_row_locks_does_not_probe_non_mysql_dialects(self):
+        query = mock.Mock()
+        bind, engine = self._mysql_bind("irrelevant")
+        bind.dialect.name = engine.dialect.name = "postgresql"
+        session = mock.Mock()
+        session.bind = bind
+        session.get_bind.return_value = bind
+
+        with mock.patch("airflow.utils.sqlalchemy.USE_ROW_LEVEL_LOCKING", True):
+            with_row_locks(query=query, session=session, skip_locked=True)
+
+        engine.connect.assert_not_called()
+        query.with_for_update.assert_called_once_with(skip_locked=True, key_share=True)
+
+    def test_with_row_locks_keeps_skip_locked_when_version_probe_fails(self):
+        query = mock.Mock()
+        bind, engine = self._mysql_bind("8.4.0")
+        engine.connect.side_effect = RuntimeError("connection gone")
+        session = mock.Mock()
+        session.bind = bind
+        session.get_bind.return_value = bind
+
+        with mock.patch("airflow.utils.sqlalchemy.USE_ROW_LEVEL_LOCKING", True):
+            with_row_locks(query=query, session=session, skip_locked=True)
+
+        query.with_for_update.assert_called_once_with(skip_locked=True, key_share=True)
+
     def test_prohibit_commit(self):
         with prohibit_commit(self.session) as guard:
             self.session.execute(text("SELECT 1"))