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

[1fanwang]

Airflow claims work with SELECT ... FOR UPDATE ... SKIP LOCKED in eleven places. On TiDB, any of
those queries that also joins another table fails outright with error 1105, which takes down
SchedulerJob._run_scheduler_loop. Nothing is scheduled until the loop is restarted.

Exception when executing SchedulerJob._run_scheduler_loop
(1105, "Can't find column airflow.dag_run.id in schema Column: [airflow.task_instance.id, ...]")
[SQL: SELECT task_instance.id, ... FROM task_instance
      INNER JOIN job ON job.id = task_instance.queued_by_job_id
      INNER JOIN dag_run ON dag_run.dag_id = task_instance.dag_id ...
      FOR UPDATE OF task_instance SKIP LOCKED]

The cause is in TiDB's planner.
LogicalLock.PruneColumns
appends every locked table's handle column to parentUsedCols so those columns survive column
pruning — but only when the lock type is supported. SKIP LOCKED is absent from
isSelectForUpdateLockType,
so the function returns early and the handle columns get pruned away while TblID2Handle still
references them. Reported as pingcap/tidb#67715
(open, sig/planner, severity/major, affects-8.5). On a query with no join there is no second
handle to lose, so the clause is instead accepted and silently ignored — the same absence from the
supported set, a different symptom.

with_row_locks() already degrades gracefully for MySQL-family servers that cannot lock at all.
This extends that to a server that accepts SKIP LOCKED and does not implement it: read the version
banner once per engine, and fall back to plain blocking FOR UPDATE. Claimers serialize instead of
skipping ahead, which is correct and only applies to servers that were never providing the
guarantee. If the version probe fails the previous behaviour is kept. PostgreSQL and MySQL are
untouched.

Testing Done

Two real airflow scheduler processes against one TiDB v8.5.1, asset-triggered Dags, 6 producer
runs. Run twice against the same deployment, once per branch:

	scheduler-loop crashes	error 1105	dag_run	task_instance
without this change	1	4	{queued: 6}	{None: 36}
with this change	0	0	{failed: 6}	{failed: 36}

Without it the loop dies and nothing is ever scheduled — all 36 task instances stay None. With it
the loop survives and all 36 are scheduled and dispatched. They then fail because this harness runs
no api-server, which Airflow 3 needs to execute tasks; the signal here is loop survival and whether
task instances are scheduled at all.

Raw logs

Reduced to two tables, no Airflow. MySQL 8.4.11 accepts both forms; TiDB rejects the second on
both v8.5.1 and current master (b76bfbc):

CREATE TABLE parent (id INT PRIMARY KEY);
CREATE TABLE child  (id INT PRIMARY KEY, parent_id INT);
INSERT INTO parent VALUES (1); INSERT INTO child VALUES (1,1);

SELECT child.id FROM child JOIN parent ON parent.id=child.parent_id FOR UPDATE OF child;
SELECT child.id FROM child JOIN parent ON parent.id=child.parent_id FOR UPDATE OF child SKIP LOCKED;

MySQL 8.4.11    FOR UPDATE OF child             -> OK
                FOR UPDATE OF child SKIP LOCKED -> OK
TiDB v8.5.1     FOR UPDATE OF child             -> OK
                FOR UPDATE OF child SKIP LOCKED -> ERROR 1105 "Can't find column bugrepro.parent.id
                                                   in schema Column: [bugrepro.child.id]"
TiDB master     FOR UPDATE OF child SKIP LOCKED -> same

Two schedulers, without this change — the loop dies:

loop crashes : 1
error 1105   : 4
dag_run      : {'queued': 6}
task_instance: {None: 36}

Two schedulers, with this change:

loop crashes : 0
error 1105   : 0
guard fired  : 1
dag_run      : {'failed': 6}
task_instance: {'failed': 36}

[warning] Database server reports as '8.0.11-tidb-v8.5.1', which accepts SKIP LOCKED
but does not honor it. Falling back to plain FOR UPDATE so concurrent schedulers
cannot claim the same rows. [airflow.utils.sqlalchemy]

Emitted SQL per backend:

backend     emits SKIP LOCKED  locking clause
----------------------------------------------------------------------------
tidb                    False  LIMIT 512 FOR UPDATE OF task_instance
mysql                    True  LIMIT 512 FOR UPDATE OF task_instance SKIP LOCKED
postgres                 True  LIMIT 512 FOR NO KEY UPDATE OF task_instance SKIP LOCKED

New unit tests against the unpatched source — they fail, which is the point:

E   Expected: with_for_update(key_share=True)
E   Actual: with_for_update(skip_locked=True, key_share=True)
FAILED ...[tidb-ignores-skip-locked]
FAILED ...[tidb-lowercase-banner]
2 failed, 3 passed

With the change: 35 passed in tests/unit/utils/test_sqlalchemy.py.

Regressions: test_scheduler_job.py -k "critical_section or executable_task_instances or row_lock or pool" — 47 passed.

Open question on the shape of this fix

This detects the engine by version banner and degrades silently. That is the weakest of three
options and I would rather land the right one than this one.

The precedent in with_row_locks() pairs its MariaDB degradation with an explicit statement that
HA scheduling is not supported there. This change does not do that: it makes an engine that cannot
provide the locking semantics Airflow needs look like it works. On the crash path in particular,
one could argue the error is the correct outcome, since it surfaces an unsupported backend
immediately rather than running on with different locking behaviour.

Alternatives, in increasing order of how well they generalise:

1. Docs only. State the capability contract - Airflow needs an engine that actually implements
   SKIP LOCKED and NOWAIT, and engines that do not are unsupported for HA scheduling. No code,
   covers every fork.
2. Capability probe at startup, in airflow db check. Run a two-session probe once and fail
   fast with an actionable message instead of crashing in the scheduler loop later. Detects
   behaviour rather than matching names, so it covers forks nobody has heard of.
3. This PR. A per-query banner denylist. Needs a new entry per engine and hides the problem.

Happy to convert this to (1) or (2) if that is the preferred direction.

Note: the main scheduler claim in _executable_task_instances_to_queued runs inside the
slot_pool ... FOR UPDATE NOWAIT critical section, and TiDB implements NOWAIT correctly, so
concurrent schedulers already serialize there. This change is about the claim queries that rely on
SKIP LOCKED alone.

---

Was generative AI tooling used to co-author this PR?

• Yes — GitHub Copilot CLI (Claude Opus 5)

Generated-by: GitHub Copilot CLI (Claude Opus 5) following the guidelines

[vikramkoka]

Oh, fascinating

Stefan, I am not familiar with TiDB.
Doesn't have to be in response to this PR, but curious about your thoughts on why this for the Airflow meta database?
Presumably for scaling, but more details would be very useful.

[1fanwang]

Oh, fascinating

Stefan, I am not familiar with TiDB. Doesn't have to be in response to this PR, but curious about your thoughts on why this for the Airflow meta database? Presumably for scaling, but more details would be very useful.

Hey Vikram, I was actually just reading more about #46175 and #65453.

To be forthcoming - I'm not proposing official support for Airflow on a TiDB backend, and this PR isn't meant to be a step toward asking for it, at least for now. Let me share

1. what I'm doing internally (Since some of this info is already public 1, 2
2. and what I think is useful to the community today:

For scale context, our largest single cluster is reaching 25k+ Dags and still growing. Most of the scaling problems we've hit have answers that stay close to upstream: add schedulers, tune the executor, etc.. The metadata DB is the one that doesn't really scale other than vertically, *can't easily be sharded, since the scheduler critical section, TI state, XCom and event logs all write to one primary, and we always just get a bigger box (we've done that many times) or changes that drift away from OSS Airflow. I'd rather not drift.

We've also been trying read/write splitting to take some pressure off from it, since we already run read replicas. Either route reads explicitly in Airflow's own source, or put query routing rules in a ProxySQL layer so Airflow core can stay generic. No numbers to share yet, but the challenge is clear - it only moves read load, and the scheduler's hot path is writes/txns. That makes distributing writes interesting, and TiDB uses the MySQL wire protocol. I haven't reached to the point to benchmarked it yet, so "scales writes" is a motivation and not a result I can show at this point.

So far this is just a small local cluster with Airflow pointed at it, checking the SQL queries (the ones scheduler depends on). Most of it holds up: pessimistic FOR UPDATE blocks, NOWAIT errors, GET_LOCK is exclusive, READ COMMITTED behaves, savepoints roll back, FK cascades are enforced. The 3 prs I opened addresses some of the minor issues I found during testing.

For this PR tho, my original thought is it isn't really a TiDB-specific problem. Any server that accepts SKIP LOCKED and quietly drops it hands two schedulers the same rows, and Airflow never finds out (it should fail loudly instead). with_row_locks() already degrades gracefully when a MySQL-family engine can't lock at all.

I do plan on trying this on our internal Airflow and TiDB clusters. Happy to share what we find running Airflow on TiDB at scale if folks are interested.

The three that came out of the exercise, for reference:

• Allow configuring the collation of asset name columns #71249
• Do not send SKIP LOCKED to servers that do not implement it #71250
• Fix MySQL metadata database failing when only PyMySQL is installed #71251

[uranusjr]

I don’t quite understand. So TiDB silently ignores SKIP LOCKED, but this PR simply makes Airflow not send that. So the end result is unchanged? Why is this PR needed?