Lee-W commented on code in PR #62501:
URL: https://github.com/apache/airflow/pull/62501#discussion_r3644880841
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -518,18 +548,19 @@ def _queue_dagruns(
# mapped) tasks update the same asset, this can fail with a unique
# constraint violation.
#
- # Where the dialect supports a single-statement "insert, ignore on
+ # Where the dialect supports a single-statement "insert, update on
# conflict" we use it; it is atomic, avoids the per-row SAVEPOINT
churn,
# and holds locks for far less time (which on MySQL/InnoDB also makes
the
# concurrent fan-out much less deadlock-prone). Otherwise we
"fallback" to
# a nested transaction per row. Either way the rows are added in the
same
# transaction where `ti.state` is changed.
dialect_name = get_dialect_name(session)
- if dialect_name == "postgresql":
Review Comment:
Let's do a
```python
if TYPE_CHECKING:
assert dialect_name is not None
```
here. so we don't need to add None to
`_queue_dagruns_nonpartitioned_conflict_update`
##########
airflow-core/tests/unit/assets/test_manager.py:
##########
@@ -234,18 +235,23 @@ def test_queue_dagruns_routes_by_dialect(self,
dialect_name, expected_helper):
dags_to_queue={dag},
partition_key=None,
partition_date=None,
- event=mock.MagicMock(),
+ event=event,
task_instance=None,
session=session,
)
- mock_helper.assert_called_once_with(1, {dag}, session)
+ if expected_helper == "_queue_dagruns_nonpartitioned_conflict_update":
+ mock_helper.assert_called_once_with(1, {dag}, event, session,
dialect_name)
+ else:
+ mock_helper.assert_called_once_with(1, {dag}, event, session)
Review Comment:
```suggestion
elif expected_helper == "_queue_dagruns_nonpartitioned_mysql":
mock_helper.assert_called_once_with(1, {dag}, event, session)
else:
raise AssertionError(f"Unexpected expected_helper:
{expected_helper}")
```
for explicity
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -717,14 +758,20 @@ def _get_or_create_apdr(
@classmethod
def _queue_dagruns_nonpartitioned_slow_path(
- cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+ cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent,
session: Session
) -> None:
def _queue_dagrun_if_needed(dag: DagModel) -> str | None:
- item = AssetDagRunQueue(target_dag_id=dag.dag_id,
asset_id=asset_id)
+ item = AssetDagRunQueue(target_dag_id=dag.dag_id,
asset_id=asset_id, created_at=event.timestamp)
# Don't error whole transaction when a single RunQueue item
conflicts.
#
https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#using-savepoint
try:
with session.begin_nested():
+ existing = session.get(
+ AssetDagRunQueue, {"target_dag_id": dag.dag_id,
"asset_id": asset_id}
+ )
Review Comment:
I think we still need this?
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -3875,6 +3875,38 @@ def
test_runtime_partition_key_does_not_backfill_dag_run_when_none(dag_maker, se
assert dr.partition_key is None
[email protected]("sqlite")
+def test_runtime_partition_key_backfill_does_not_deadlock_on_sqlite(dag_maker,
session):
+ """Regression test for the SQLite ``database is locked`` deadlock between
the
+ writes in ``register_asset_changes_in_db`` and the second connection that
+ ``_create_asset_event`` used to
+ open. On file-based SQLite (the default ``-b sqlite`` test backend) the two
+ connections compete for the same RESERVED lock; the SQLite branch of
+ ``_create_asset_event`` must add the event directly to the caller's session
+ instead of opening a second connection.
+ """
Review Comment:
```suggestion
"""Regression test for the SQLite ``database is locked`` deadlock
between the
writes in ``register_asset_changes_in_db`` and the second connection that
``_create_asset_event`` used to open.
On file-based SQLite (the default ``-b sqlite`` test backend) the two
connections compete for the same RESERVED lock; the SQLite branch of
``_create_asset_event`` must add the event directly to the caller's
session
instead of opening a second connection.
```
the breakline was hard to read
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -3875,6 +3875,38 @@ def
test_runtime_partition_key_does_not_backfill_dag_run_when_none(dag_maker, se
assert dr.partition_key is None
[email protected]("sqlite")
+def test_runtime_partition_key_backfill_does_not_deadlock_on_sqlite(dag_maker,
session):
+ """Regression test for the SQLite ``database is locked`` deadlock between
the
+ writes in ``register_asset_changes_in_db`` and the second connection that
+ ``_create_asset_event`` used to
+ open. On file-based SQLite (the default ``-b sqlite`` test backend) the two
+ connections compete for the same RESERVED lock; the SQLite branch of
+ ``_create_asset_event`` must add the event directly to the caller's session
+ instead of opening a second connection.
+ """
+ asset = Asset(name="hello")
+ with dag_maker(dag_id="rt_pk_backfill_sqlite",
schedule=PartitionedAtRuntime()) as dag:
+ EmptyOperator(task_id="hi", outlets=[asset])
+ dr = dag_maker.create_dagrun(session=session)
+ assert dr.partition_key is None
+ [ti] = dr.get_task_instances(session=session)
+
+ # Must not raise sqlite3.OperationalError: database is locked.
Review Comment:
```suggestion
# Must not raise "sqlite3.OperationalError: database is locked."
```
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -789,33 +819,46 @@ def _get_or_create_apdr(
return apdr
@classmethod
- def _queue_dagruns_nonpartitioned_slow_path(
- cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+ def _queue_dagruns_nonpartitioned_mysql(
+ cls, asset_id: int, dags_to_queue: set[DagModel], event: AssetEvent,
session: Session
) -> None:
- def _queue_dagrun_if_needed(dag: DagModel) -> str | None:
- item = AssetDagRunQueue(target_dag_id=dag.dag_id,
asset_id=asset_id)
- # Don't error whole transaction when a single RunQueue item
conflicts.
- #
https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#using-savepoint
- try:
- with session.begin_nested():
- session.merge(item)
- except exc.IntegrityError:
- cls.logger().debug("Skipping record %s", item, exc_info=True)
- return dag.dag_id
+ from sqlalchemy import case
+ from sqlalchemy.dialects.mysql import insert
- queued_results = (_queue_dagrun_if_needed(dag) for dag in
dags_to_queue)
- if queued_dag_ids := [r for r in queued_results if r is not None]:
- cls.logger().debug("consuming dag ids %s", queued_dag_ids)
+ values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue]
+ stmt = insert(AssetDagRunQueue).values(asset_id=asset_id,
created_at=event.timestamp)
+
+ update_stmt = stmt.on_duplicate_key_update(
+ created_at=case(
+ (stmt.inserted.created_at >= AssetDagRunQueue.created_at,
stmt.inserted.created_at),
+ else_=AssetDagRunQueue.created_at,
+ )
+ )
+ session.execute(update_stmt, values)
@classmethod
- def _queue_dagruns_nonpartitioned_postgres(
- cls, asset_id: int, dags_to_queue: set[DagModel], session: Session
+ def _queue_dagruns_nonpartitioned_conflict_update(
+ cls,
+ asset_id: int,
+ dags_to_queue: set[DagModel],
+ event: AssetEvent,
+ session: Session,
+ dialect_name: str | None,
Review Comment:
https://github.com/apache/airflow/pull/62501/changes#r3644880841
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
