diff --git a/airflow-core/newsfragments/71250.bugfix.rst b/airflow-core/newsfragments/71250.bugfix.rst new file mode 100644 index 0000000000000..44d0bf36da14a --- /dev/null +++ b/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. diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index 61ac586bda6cf..32ce3dc9e9b5c 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/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 diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 8dbcd41313055..ed2b8bb17a03f 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/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"))