jason810496 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3610221272
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -3676,6 +3682,88 @@ def _remove_unreferenced_triggers(self, *, session:
Session = NEW_SESSION) -> No
.execution_options(synchronize_session="fetch")
)
+ @provide_session
+ def _drain_asset_event_queue(self, *, session: Session = NEW_SESSION) ->
None:
+ """
+ Process pending asset-event registrations enqueued by the execution
API.
+
+ The task-success path enqueues an
:class:`~airflow.models.asset.AssetEventQueue` row when
+ synchronous asset-event registration is skipped or fails under
database lock contention.
+ Here we claim pending rows one at a time (``FOR UPDATE SKIP LOCKED``,
so multiple schedulers
+ do not collide), (re)run ``register_asset_changes_in_db`` to create
the ``AssetEvent`` and
+ ``AssetDagRunQueue`` rows, and delete the queue row in the same
transaction so it is
+ processed exactly once. A row that keeps failing is retried on later
passes until it exceeds
+ ``asset_event_queue_max_attempts``, after which it is left in place
(parked) and surfaced via
+ the ``asset.event_queue.failures`` metric for an operator to inspect.
+ """
+ from airflow.api_fastapi.execution_api.datamodels.asset import
AssetProfile
+
+ batch_size = conf.getint("scheduler", "asset_event_queue_batch_size",
fallback=100)
+ max_attempts = conf.getint("scheduler",
"asset_event_queue_max_attempts", fallback=5)
+
+ processed = 0
+ # Track rows attempted this pass so a freshly-failed row (whose
attempts were just bumped)
+ # is not re-claimed and hammered within the same drain; it waits for
the next pass instead.
+ seen: set[UUID] = set()
+ for _ in range(batch_size):
+ conditions = [AssetEventQueue.attempts < max_attempts]
+ if seen:
+ conditions.append(AssetEventQueue.ti_id.notin_(seen))
+ row = session.scalars(
+ with_row_locks(
+
select(AssetEventQueue).where(*conditions).order_by(AssetEventQueue.created_at).limit(1),
+ of=AssetEventQueue,
+ session=session,
+ skip_locked=True,
+ )
+ ).first()
+ if row is None:
+ break
+
+ ti_id = row.ti_id
+ seen.add(ti_id)
+ ti = session.get(TaskInstance, ti_id)
+ try:
Review Comment:
We should fix the case that the TI being clear as we set the FK constraint
between `AssetEventQueue` and `TaskInstance`.
When every "TaskInstance clear", we will set the new `ti_id` UUID. If the
`AssetEventQueue` records are not yet consumed up, we will enter the `if ti is
not None:` path and drop the AssetEvent ingestion.
https://github.com/apache/airflow/blob/ffde96203e76c0c27ccf1ae6eeaad270653248bf/airflow-core/src/airflow/models/taskinstance.py#L1034-L1040
---
To fix the above scenario, we should update the `AssetEventQueue.ti_id` to
new ti_id in the same tx.
```python
def prepare_db_for_next_try(self, session: Session):
from airflow.models.taskinstancehistory import TaskInstanceHistory
TaskInstanceHistory.record_ti(self, session=session)
session.execute(delete(TaskReschedule).filter_by(ti_id=self.id))
session.execute(
update(AssetEventQueue).where(AssetEventQueue.ti_id ==
self.id).values(ti_id=new_id)
) # new
self.id = uuid7()
```
--
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]
