fix(asset-parition): add a mutex lock to AssetModel when updating APDR by Lee-W · Pull Request #59183 · apache/airflow

Lee-W · 2025-12-08T03:54:16Z

Why

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W · 2025-12-08T03:54:16Z

Why

closes: #58919

What

Add a row lock to AssetModel when updating APDR for MySQL and PostgreSQL. Since SQLIte does not support row lock, add retry mechanism if OperationError with Database is locked is encountered.

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-12-08T23:17:15Z

I don't think it's a good solution

SELECT FOR UPDATE will not lock anything when ADPR lock is missing.

If the task switch happens after ROW lock in first event handling, both select for updates will run and return no row, and this will create both ADPRs:

Time	Event 1	Event 2
0	SELECT FOR UPDATE ADR (NO ROW FOUND)
1		SELECT FOR UPDATE ADR (NO ROW FOUND)
2		CREATE ADPR
3	CREATE ADPR

Lee-W · 2025-12-09T10:40:36Z

Just updated it to lock AssetModel inatead of APDR

potiuk

Way better I think. @dstandish ?

potiuk

Hmm. took a closer look and I have some improvement proposal and some surprise when looking at the test.

potiuk

Much better - htough I think maybe we should run the test few times if it's quick? Lack of the race condition does not say it cannot occur if it's only executed once :)

I did not know SQLite does not support row-level locks. Good solution with extra table.

potiuk · 2025-12-10T19:53:24Z

I think you need to rebase to get rid of the issues that were fixed in main already

dstandish

left some suggestions

dstandish

small question but looks ok

…te mutex lock

…Lite only)

apache#59183)

Lee-W force-pushed the resolve-apdr-race-condition branch from 4c835f4 to edd36e0Compare December 8, 2025 08:09

Lee-W mentioned this pull request Dec 8, 2025
Resolve race condition re creation of APDR records #58919
Closed
1 task

Lee-W force-pushed the resolve-apdr-race-condition branch 2 times, most recently from 95e4053 to c224bccCompare December 8, 2025 14:23

Lee-W marked this pull request as ready for review December 8, 2025 14:33

Lee-W requested a review from dstandish December 8, 2025 14:33

Lee-W changed the title ~~fix(asset-parition): add a mutex lock to adpr updating~~fix(asset-parition): add a mutex lock to APDR updatingDec 8, 2025

Lee-W marked this pull request as draft December 9, 2025 01:18

Lee-W force-pushed the resolve-apdr-race-condition branch 2 times, most recently from 386669a to 346eba2Compare December 9, 2025 06:00

jason810496 self-requested a review December 9, 2025 06:53

Lee-W changed the title ~~fix(asset-parition): add a mutex lock to APDR updating~~fix(asset-parition): add a mutex lock to AssetModel when updating APDRDec 9, 2025

Lee-W marked this pull request as ready for review December 9, 2025 07:42

potiuk approved these changes Dec 9, 2025
View reviewed changes

potiuk reviewed Dec 9, 2025
View reviewed changes

Comment threadairflow-core/src/airflow/assets/manager.py Outdated

potiuk reviewed Dec 9, 2025
View reviewed changes

Comment threadairflow-core/tests/unit/assets/test_manager.py Outdated

potiuk requested changes Dec 9, 2025
View reviewed changes

Lee-W force-pushed the resolve-apdr-race-condition branch 2 times, most recently from 04023ee to 9688202Compare December 10, 2025 09:05

Lee-W requested review from XD-DENG, ashb and ephraimbuddy as code owners December 10, 2025 09:05

Lee-W requested a review from potiuk December 10, 2025 09:49

Lee-W force-pushed the resolve-apdr-race-condition branch from 9688202 to 11ba716Compare December 10, 2025 10:05

potiuk approved these changes Dec 10, 2025
View reviewed changes

potiuk reviewed Dec 10, 2025
View reviewed changes

Comment threadairflow-core/src/airflow/assets/manager.py Outdated

Lee-W force-pushed the resolve-apdr-race-condition branch from 5892a87 to 93f84f4Compare December 11, 2025 01:53

dstandish reviewed Dec 16, 2025
View reviewed changes

Lee-W force-pushed the resolve-apdr-race-condition branch 3 times, most recently from 67adb42 to e7d33fbCompare December 17, 2025 09:40

dstandish reviewed Dec 17, 2025
View reviewed changes

Comment threadairflow-core/src/airflow/assets/manager.py Outdated

dstandish approved these changes Dec 17, 2025
View reviewed changes

Lee-W force-pushed the resolve-apdr-race-condition branch from e7d33fb to 8c320eeCompare December 18, 2025 07:57

Lee-W added 14 commits December 18, 2025 17:01

fix(asset-parition): add a mutex lock to adpr updating
1cadb01

fix(asset-partition): fix the sqlite case
4895a15

fix(asset-partition): lock asset row instead
c768693

feat(asset-parition): add asset_partition_dag_run_mutex_lock for sqli…
ef31619
…te mutex lock

refactor(asset-partition): add APDRMutexLock table for all db
33b6af2

feat(asset-partition): add key_share to row lock
5946520

refactor: remove unnecessay AssetPartitionDagRunMutexLock
6b1411a

test(assets): improve test_get_or_create_apdr_race_condition
a7ec1e3

refactor: rename _acquire_apdr_lock as _lock_asset_model
41b05b5

refactor: remove unused variables
775c3e9

fix(assets): change asset_model lock retry delay from 0.05 to 0.1 (SQ…
2d5ac9f
…Lite only)

refactor: rewrite the try logic of SQLite handling
e699f04

refactor(asset-partition): remove unnecessary session flush
e908e9d

fixup! refactor(asset-partition): remove unnecessary session flush
a418c89

Lee-W force-pushed the resolve-apdr-race-condition branch from 8c320ee to a418c89Compare December 18, 2025 09:01

Lee-W merged commit c36e5ec into apache:mainDec 18, 2025
329 of 333 checks passed

Lee-W deleted the resolve-apdr-race-condition branch December 18, 2025 14:20

Lohith625 pushed a commit to Lohith625/airflow that referenced this pull request Dec 19, 2025
fix(asset-parition): add a mutex lock to AssetModel when updating APDR (
b3891b7
apache#59183)

stegololz pushed a commit to stegololz/airflow that referenced this pull request Jan 9, 2026
fix(asset-parition): add a mutex lock to AssetModel when updating APDR (
3335e7c
apache#59183)

jhgoebbert pushed a commit to jhgoebbert/airflow_Owen-CH-Leung that referenced this pull request Feb 8, 2026
fix(asset-parition): add a mutex lock to AssetModel when updating APDR (
09d203d
apache#59183)

Subham-KRLX pushed a commit to Subham-KRLX/airflow that referenced this pull request Mar 4, 2026
fix(asset-parition): add a mutex lock to AssetModel when updating APDR (
361b4f9
apache#59183)

Mercury2699 mentioned this pull request Jun 4, 2026
Duplicate DAG runs for PartitionedAssetTimetable with multiple schedulers (HA) — race not covered by #60773 #68045
Closed

bujjibabukatta mentioned this pull request Jun 5, 2026
fix(scheduler): add row lock to AssetPartitionDagRun fetch to prevent… #68061
Closed