kaxil commented on code in PR #75943:
URL: https://github.com/apache/airflow/pull/75943#discussion_r4410980688
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2945,42 +2972,45 @@ def _select_consumed_asset_events(
AssetAliasModel.scheduled_dags.any(DagScheduleAssetAliasReference.dag_id ==
dag.dag_id)
),
)
- return list(
- session.scalars(
- select(AssetEvent)
- .where(
- event_predicate,
- ~(
- select(association_table.c.event_id)
- .join(DagRun, DagRun.id ==
association_table.c.dag_run_id)
- .where(
- DagRun.dag_id == dag.dag_id,
- association_table.c.event_id == AssetEvent.id,
- )
- .exists()
- ),
- )
- .order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc())
- .limit(limit)
+ return session.scalars(
+ select(AssetEvent)
+ .where(
+ event_predicate,
+ ~_consumed_by_dag(dag.dag_id, AssetEvent.id),
)
- )
+ .order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc())
+ .limit(limit)
+ ).all()
def _delete_consumed_asset_records(
- self, *, records: Sequence[AssetDagRunQueue], dag_id: str, session:
Session
+ self,
+ *,
+ dag_id: str,
+ max_event_id: int,
+ keep_event_id: int | None = None,
+ session: Session,
) -> None:
- # Delete only consumed ADRQ rows to avoid dropping newly queued events
- # (e.g. DagRun triggered by asset A while a new event for asset B
arrives).
- result = cast(
- "CursorResult",
- session.execute(
- delete(AssetDagRunQueue).where(
- tuple_(
- AssetDagRunQueue.target_dag_id,
- AssetDagRunQueue.asset_event_id,
- ).in_((record.target_dag_id, record.asset_event_id) for
record in records)
- )
- ),
+ """
+ Delete ADRQ rows whose events have been consumed by a run of the Dag.
+
+ Only consumed rows are deleted to avoid dropping newly queued events.
+ Rows are matched against the events consumed by the Dag's runs (i.e.
+ ``DagRun.consumed_asset_events``) with a subquery, so the caller must
+ flush newly consumed events before calling this.
+
+ :param max_event_id: Largest event ID among the ADRQ rows locked by
this
+ scheduler. Rows queued after the lock was taken are not touched.
+ :param keep_event_id: If given, the ADRQ row of this event is kept
even if
+ the event is consumed.
+ """
+ stmt = delete(AssetDagRunQueue).where(
+ AssetDagRunQueue.target_dag_id == dag_id,
+ AssetDagRunQueue.asset_event_id <= max_event_id,
Review Comment:
This range now also matches rows SKIP LOCKED skipped, and a DELETE waits on
a locked row instead of skipping it. If one of the queuedEvents DELETE
endpoints has locked row A and is waiting on a row this scheduler holds, this
statement reaching A closes the cycle. A Postgres repro with the timing forced
by a sleep gets `DeadlockDetected` at this head, while on the base commit the
API call just waits. The description says the scheduler's retry recovers, but
Postgres aborts whichever backend runs the deadlock check first, which can be
the API request, and that request gets a 500 with no retry. Could the DELETE be
limited to rows this transaction locked, e.g. `asset_event_id IN (SELECT ...
WHERE target_dag_id = :d AND asset_event_id <= :max FOR UPDATE SKIP LOCKED)`?
The parameter count stays constant, though MySQL needs a derived table for the
self-reference.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2926,13 +2939,27 @@ def _select_consumed_asset_events(
self,
*,
dag: SerializedDAG,
- records: Sequence[AssetDagRunQueue],
+ max_event_id: int,
limit: int | None = None,
session: Session,
- ) -> list[AssetEvent]:
- """Select unconsumed events referenced by a Dag's locked ADRQ rows."""
- referenced_event_ids = {record.asset_event_id for record in records}
- event_predicate: ColumnElement[bool] =
AssetEvent.id.in_(referenced_event_ids)
+ ) -> Sequence[AssetEvent]:
+ """
+ Select unconsumed events referenced by a Dag's locked ADRQ rows.
+
+ *max_event_id* is the largest event ID among the rows locked by this
+ scheduler, so events queued after the lock was taken are left for the
Review Comment:
I don't think the bound guarantees this. With `catchup=True` the two OR'd
arms below have no id bound, so they still pick up events queued after the
lock. And event ids aren't assigned in commit order, so on Postgres a lower id
can commit after the lock and still fall inside the bound. Both behave the same
as the old code, so it's only the docstring (and the `max_event_id` param doc
on `_delete_consumed_asset_records`) that overpromises.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -13034,6 +13038,190 @@ def
test_create_dagruns_asset_and_time_late_arrival_uses_oldest_pending_logical_
)
[email protected]
+def _capture_max_bind_params(session: Session):
+ """Track the largest number of bind parameters sent in a single SQL
statement."""
+ counts: list[int] = [0]
+
+ def _listener(conn, cursor, statement, parameters, context, executemany):
+ if not executemany and parameters:
+ counts.append(len(parameters))
+
+ engine = session.get_bind()
+ sa_event.listen(engine, "before_cursor_execute", _listener)
+ try:
+ yield counts
+ finally:
+ sa_event.remove(engine, "before_cursor_execute", _listener)
+
+
+# Large enough that a literal IN list of queued events dwarfs any fixed-size
query.
+_MANY_QUEUED_EVENTS = 300
+
+
+def _queue_many_asset_events(
+ *,
+ dag_id: str,
+ asset_ids: list[int],
+ timestamp: datetime.datetime,
+ session: Session,
+) -> None:
+ events = [
+ AssetEvent(asset_id=asset_id, timestamp=timestamp)
+ for asset_id in asset_ids
+ for _ in range(_MANY_QUEUED_EVENTS)
+ ]
+ session.add_all(events)
+ session.flush()
+ session.add_all(
+ AssetDagRunQueue(asset_id=e.asset_id, target_dag_id=dag_id,
asset_event_id=e.id) for e in events
+ )
+ session.flush()
+
+
[email protected]_serialized_dag
[email protected]("batch_asset_events", [False, True])
+def test_create_dagruns_asset_triggered_many_queued_events_no_literal_in(
+ session: Session, dag_maker, batch_asset_events: bool
+):
+ """Queued events must not be sent to the database as a literal list
(#75865)."""
Review Comment:
Could we drop the `(#75865)` here and in the docstring on line 13134? The
test name and docstring already describe the behaviour, and the PR links the
issue.
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -7302,65 +7302,69 @@ def
test_create_dag_runs_asset_triggered_skips_stale_triggered_date(self, sessio
assert session.scalars(select(DagRun).where(DagRun.dag_id ==
dag_model.dag_id)).one_or_none() is None
@pytest.mark.need_serialized_dag
+ @pytest.mark.parametrize(
+ ("locked_index", "expected_consumed_indexes",
"expected_queued_indexes"),
+ [
+ # The skipped row has a higher event ID than any locked row, so it
is
+ # outside the bound and left queued for the next scheduling loop.
+ pytest.param(0, {0}, {1}, id="skipped-row-newer"),
+ # The skipped row has a lower event ID than the locked row, so it
is
+ # within the bound, and its event is consumed and dequeued as well.
+ pytest.param(1, {0, 1}, set(), id="skipped-row-older"),
+ ],
+ )
def test_create_dag_runs_asset_triggered_deletes_only_selected_adrq_rows(
Review Comment:
In the `skipped-row-older` case the skipped row is consumed and dequeued
too, so "deletes only selected" no longer describes what this asserts. Maybe
rename it to something like `..._consumes_skipped_rows_within_bound`?
--
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]
