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 that also joins another table fails with error 1105 and takes down
SchedulerJob._run_scheduler_loop. Nothing is scheduled until it restarts. One scheduler is enough
to hit it.

Exception when executing SchedulerJob._run_scheduler_loop
(1105, "Can't find column airflow.dag_run.id in schema Column: [airflow.task_instance.id, ...]")

LogicalLock.PruneColumns
protects each locked table's handle column from pruning, but only for supported lock types.
SKIP LOCKED is missing from
isSelectForUpdateLockType,
so the handle is pruned while TblID2Handle still references it. Filed as
pingcap/tidb#67715. Without a join there is no
second handle to lose, so the clause is silently ignored instead.

with_row_locks() already degrades for MySQL-family servers that cannot lock. This extends it to
servers that accept SKIP LOCKED without implementing it: read the banner once per engine, fall
back to plain FOR UPDATE. If the probe fails, previous behaviour is kept. PostgreSQL and MySQL
are untouched.

Testing Done

An api-server and two airflow scheduler processes against one TiDB v8.5.1. Six producer Dag runs
fanning out to six asset-scheduled consumers. Same deployment, once per branch:

	loop crashes	error 1105	dag_run	task_instance	asset events
without this change	1	4	producer queued x6	{None: 36}	0
with this change	0	0	all 7 Dags success	{success: 66}	36

Raw logs

CREATE TABLE parent (id INT PRIMARY KEY);
CREATE TABLE child  (id INT PRIMARY KEY, parent_id INT);
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

join predicate               irrelevant
FOR UPDATE OF <table>        does not restrict which handles are tracked
PK-less table                fails on _tidb_rowid
handle added to SELECT list  error vanishes  -> pruning, not locking

without:

loop crashes : 1
error 1105   : 4
dag_run      : {('asset_producer', 'queued'): 6}
task_instance: {None: 36}
asset_events : 0

with:

loop crashes : 0
error 1105   : 0
guard fired  : 1

dag_run: ('asset_consumer_0', 'success', 5)   ... consumers 1-5 identical
dag_run: ('asset_producer',   'success', 6)
task_instance: {'success': 66}
asset_events : 36
retries      : 0

[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. [airflow.utils.sqlalchemy]

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

tests/unit/utils/test_sqlalchemy.py                     35 passed (2 fail unpatched)
test_scheduler_job.py -k critical_section/row_lock/pool 47 passed

_executable_task_instances_to_queued is unaffected: it runs inside the slot_pool ... FOR UPDATE NOWAIT critical section, which TiDB implements correctly. Banner matching is the weakest available
fix — a capability probe in airflow db check, or docs stating the contract, would generalise
better. Happy to convert.

---

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.

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.

[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?

[1fanwang]

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?

Closing the PR as not applicable — sorry for the noise @uranusjr TP you are right that this changes nothing, adds no value on making Airflow run on TIDB today.

Testing Airflow on TiDB further showed it can't even keep a scheduler loop alive, I think maybe due to a TiDB planner bug pingcap/tidb#67715, will have to look into this further and test, closing for now