jason810496 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3671539907
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -497,12 +502,38 @@ def ti_update_state(
extra=json.dumps({"host_name": hostname}) if hostname else
None,
)
)
+ # Durably record the successful task's asset events in the SAME
transaction that commits the
+ # state, so the marker can never be lost in a crash between this
commit and their
+ # registration. The scheduler drain is the single writer that
registers them; the request
+ # never runs asset registration itself, which is what keeps the
task_instance row lock short
+ # under high fan-out. Only a genuine RUNNING->SUCCESS transition
reaches here (a duplicate
+ # SUCCESS->SUCCESS short-circuits earlier), so a completion is never
enqueued twice.
+ if (
+ updated_state == TaskInstanceState.SUCCESS
+ and isinstance(ti_patch_payload, TISuccessStatePayload)
+ and (ti_patch_payload.task_outlets or
ti_patch_payload.outlet_events)
+ ):
+ _enqueue_asset_events(
+ task_instance_id=task_instance_id,
+ dag_id=dag_id,
+ run_id=run_id,
+ task_id=task_id,
+ map_index=map_index,
+ task_outlets=ti_patch_payload.task_outlets,
+ outlet_events=ti_patch_payload.outlet_events,
+ session=session,
+ )
+ # Commit the state, log entry and durable queue marker together so
they land atomically and
+ # the row lock is released promptly.
+ session.commit()
Review Comment:
We shouldn't call the `commit` the `rollback` explicitly in the router.
```suggestion
```
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -521,7 +552,9 @@ def ti_update_state(
task_id=task_id,
map_index=map_index,
)
+ session.commit()
except Exception:
+ session.rollback()
Review Comment:
```suggestion
except Exception:
```
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -618,6 +651,58 @@ def _validate_outlet_event_partition_keys(outlet_events:
list[dict[str, Any]]) -
)
+def _asset_event_payload(
+ task_outlets: list[AssetProfile],
+ outlet_events: list[dict[str, Any]],
+ ti_key: dict[str, Any],
+) -> dict[str, Any]:
+ """Serialize a task's outlets, outlet events, and natural key into the
queue row's JSON payload."""
+ return {
+ "task_outlets": [outlet.model_dump(mode="json") for outlet in
task_outlets],
+ "outlet_events": outlet_events,
+ # The scheduler drain resolves the live task instance by this natural
key rather than by the
+ # surrogate ``ti_id``. Clearing a task reassigns its uuid7 id, and
relying on the foreign
+ # key's ON UPDATE CASCADE to re-point the row only works on Postgres
-- SQLite does not
+ # enforce foreign keys in Airflow's production engine -- so a
natural-key lookup is what
+ # keeps the pending events reachable on every backend.
+ "ti_key": ti_key,
Review Comment:
The comment here is too verbose. `_enqueue_asset_events` already explained
the overall archciture.
```suggestion
"ti_key": ti_key,
```
##########
airflow-core/src/airflow/models/asset.py:
##########
@@ -788,6 +789,101 @@ def __repr__(self):
return f"{self.__class__.__name__}({', '.join(args)})"
+class AssetEventQueue(Base):
+ """
+ Durable marker of asset events a successful task emitted, awaiting
registration.
+
+ On task success the execution API commits one of these rows atomically
with the task
+ state instead of registering the asset events inline, so the
``ti_update_state`` request
+ holds the ``task_instance`` row lock only for the state write plus this
insert rather than
+ for the whole ``register_asset_changes_in_db`` call (which was the source
of API-server
+ lock contention under high fan-out). The scheduler drains this table,
resolves the live
+ task instance by natural key, runs
+
:meth:`~airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db`
to create
+ the ``AssetEvent`` and ``AssetDagRunQueue`` rows, and deletes the queue
row once that write
+ commits. The in-process runner behind ``dag.test`` has no scheduler, so it
drains the row
+ itself via :func:`register_pending_asset_events` right after the task
finishes.
+
+ ``ti_id`` is the primary key: at most one pending registration exists per
task
+ instance, and the row is cascade-deleted if the task instance is removed.
+ """
+
+ ti_id: Mapped[UUID] = mapped_column(sa.Uuid(), primary_key=True,
nullable=False)
+ # Both the emitted task outlets and the outlet events live in one JSON
payload
+ # (``{"task_outlets": [...], "outlet_events": [...], "ti_key": {...}}``).
The queue is a durable
+ # buffer only ever read back in full when draining, so a single column
keeps the enqueue cheap.
+ payload: Mapped[dict] = mapped_column(sa.JSON(), nullable=False,
default=dict)
+ attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0,
server_default="0")
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
+
+ __tablename__ = "asset_event_queue"
+ __table_args__ = (
+ PrimaryKeyConstraint(ti_id, name="asset_event_queue_pkey"),
+ ForeignKeyConstraint(
+ (ti_id,),
+ ["task_instance.id"],
+ name="aeq_ti_fkey",
+ ondelete="CASCADE",
+ # Referential cleanup only. Correctness on clear does not rely on
these cascades:
+ # SQLite does not enforce foreign keys in Airflow's production
engine, so they silently
+ # no-op there. The drain re-resolves the task instance by natural
key instead, which
+ # survives the id reassignment a clear performs on every backend.
+ onupdate="CASCADE",
+ ),
+ Index("idx_asset_event_queue_created_at", created_at),
+ )
+
+ def __repr__(self):
+ return f"AssetEventQueue(ti_id={self.ti_id!r},
attempts={self.attempts!r})"
+
+
+def _register_queued_asset_event(row: AssetEventQueue, *, session: Session) ->
None:
+ """
+ Register the asset events captured in one :class:`AssetEventQueue` row,
then delete it.
+
+ Resolves the live task instance by natural key
(``dag_id``/``run_id``/``task_id``/``map_index``)
+ rather than the surrogate ``ti_id``: clearing a task reassigns its id, so
a lookup by the
+ enqueued id would miss the row on any backend that does not cascade the id
change. If the task
+ instance no longer exists there is nothing to register and the row is
simply dropped. The caller
+ owns the surrounding transaction (the scheduler wraps each row in a
savepoint; the in-process
+ runner commits the session).
+ """
+ from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
+ from airflow.models.taskinstance import TaskInstance
+
+ payload = row.payload
+ ti_key = payload["ti_key"]
+ ti = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == ti_key["dag_id"],
+ TaskInstance.run_id == ti_key["run_id"],
+ TaskInstance.task_id == ti_key["task_id"],
+ TaskInstance.map_index == ti_key["map_index"],
+ )
+ )
+ if ti is not None:
+ task_outlets = [AssetProfile.model_validate(outlet) for outlet in
payload["task_outlets"]]
+ TaskInstance.register_asset_changes_in_db(ti, task_outlets,
payload["outlet_events"], session=session)
+ session.delete(row)
+
+
+def register_pending_asset_events(*, ti_ids: Iterable[UUID], session: Session)
-> None:
Review Comment:
We should rename this method and mention that this method can only be used
for `dag.test` purpose.
This shouldn't be used at any other components.
`register_pending_asset_events` is ambigious.
##########
airflow-core/src/airflow/models/asset.py:
##########
@@ -788,6 +789,101 @@ def __repr__(self):
return f"{self.__class__.__name__}({', '.join(args)})"
+class AssetEventQueue(Base):
+ """
+ Durable marker of asset events a successful task emitted, awaiting
registration.
+
+ On task success the execution API commits one of these rows atomically
with the task
+ state instead of registering the asset events inline, so the
``ti_update_state`` request
+ holds the ``task_instance`` row lock only for the state write plus this
insert rather than
+ for the whole ``register_asset_changes_in_db`` call (which was the source
of API-server
+ lock contention under high fan-out). The scheduler drains this table,
resolves the live
+ task instance by natural key, runs
+
:meth:`~airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db`
to create
+ the ``AssetEvent`` and ``AssetDagRunQueue`` rows, and deletes the queue
row once that write
+ commits. The in-process runner behind ``dag.test`` has no scheduler, so it
drains the row
+ itself via :func:`register_pending_asset_events` right after the task
finishes.
+
+ ``ti_id`` is the primary key: at most one pending registration exists per
task
+ instance, and the row is cascade-deleted if the task instance is removed.
+ """
+
+ ti_id: Mapped[UUID] = mapped_column(sa.Uuid(), primary_key=True,
nullable=False)
+ # Both the emitted task outlets and the outlet events live in one JSON
payload
+ # (``{"task_outlets": [...], "outlet_events": [...], "ti_key": {...}}``).
The queue is a durable
+ # buffer only ever read back in full when draining, so a single column
keeps the enqueue cheap.
+ payload: Mapped[dict] = mapped_column(sa.JSON(), nullable=False,
default=dict)
+ attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0,
server_default="0")
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
+
+ __tablename__ = "asset_event_queue"
+ __table_args__ = (
+ PrimaryKeyConstraint(ti_id, name="asset_event_queue_pkey"),
+ ForeignKeyConstraint(
+ (ti_id,),
+ ["task_instance.id"],
+ name="aeq_ti_fkey",
+ ondelete="CASCADE",
+ # Referential cleanup only. Correctness on clear does not rely on
these cascades:
+ # SQLite does not enforce foreign keys in Airflow's production
engine, so they silently
+ # no-op there. The drain re-resolves the task instance by natural
key instead, which
+ # survives the id reassignment a clear performs on every backend.
+ onupdate="CASCADE",
+ ),
+ Index("idx_asset_event_queue_created_at", created_at),
+ )
+
+ def __repr__(self):
+ return f"AssetEventQueue(ti_id={self.ti_id!r},
attempts={self.attempts!r})"
+
+
+def _register_queued_asset_event(row: AssetEventQueue, *, session: Session) ->
None:
Review Comment:
This should be a public method, scheduler will reference this directly.
```suggestion
def register_queued_asset_event(row: AssetEventQueue, *, session: Session)
-> None:
```
##########
airflow-core/src/airflow/models/asset.py:
##########
@@ -788,6 +789,101 @@ def __repr__(self):
return f"{self.__class__.__name__}({', '.join(args)})"
+class AssetEventQueue(Base):
+ """
+ Durable marker of asset events a successful task emitted, awaiting
registration.
+
+ On task success the execution API commits one of these rows atomically
with the task
+ state instead of registering the asset events inline, so the
``ti_update_state`` request
+ holds the ``task_instance`` row lock only for the state write plus this
insert rather than
+ for the whole ``register_asset_changes_in_db`` call (which was the source
of API-server
+ lock contention under high fan-out). The scheduler drains this table,
resolves the live
+ task instance by natural key, runs
+
:meth:`~airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db`
to create
+ the ``AssetEvent`` and ``AssetDagRunQueue`` rows, and deletes the queue
row once that write
+ commits. The in-process runner behind ``dag.test`` has no scheduler, so it
drains the row
+ itself via :func:`register_pending_asset_events` right after the task
finishes.
+
+ ``ti_id`` is the primary key: at most one pending registration exists per
task
+ instance, and the row is cascade-deleted if the task instance is removed.
+ """
+
+ ti_id: Mapped[UUID] = mapped_column(sa.Uuid(), primary_key=True,
nullable=False)
+ # Both the emitted task outlets and the outlet events live in one JSON
payload
+ # (``{"task_outlets": [...], "outlet_events": [...], "ti_key": {...}}``).
The queue is a durable
+ # buffer only ever read back in full when draining, so a single column
keeps the enqueue cheap.
+ payload: Mapped[dict] = mapped_column(sa.JSON(), nullable=False,
default=dict)
+ attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0,
server_default="0")
+ created_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
+
+ __tablename__ = "asset_event_queue"
+ __table_args__ = (
+ PrimaryKeyConstraint(ti_id, name="asset_event_queue_pkey"),
+ ForeignKeyConstraint(
+ (ti_id,),
+ ["task_instance.id"],
+ name="aeq_ti_fkey",
+ ondelete="CASCADE",
+ # Referential cleanup only. Correctness on clear does not rely on
these cascades:
+ # SQLite does not enforce foreign keys in Airflow's production
engine, so they silently
+ # no-op there. The drain re-resolves the task instance by natural
key instead, which
+ # survives the id reassignment a clear performs on every backend.
+ onupdate="CASCADE",
+ ),
+ Index("idx_asset_event_queue_created_at", created_at),
+ )
+
+ def __repr__(self):
+ return f"AssetEventQueue(ti_id={self.ti_id!r},
attempts={self.attempts!r})"
+
+
+def _register_queued_asset_event(row: AssetEventQueue, *, session: Session) ->
None:
+ """
+ Register the asset events captured in one :class:`AssetEventQueue` row,
then delete it.
+
+ Resolves the live task instance by natural key
(``dag_id``/``run_id``/``task_id``/``map_index``)
+ rather than the surrogate ``ti_id``: clearing a task reassigns its id, so
a lookup by the
+ enqueued id would miss the row on any backend that does not cascade the id
change. If the task
+ instance no longer exists there is nothing to register and the row is
simply dropped. The caller
+ owns the surrounding transaction (the scheduler wraps each row in a
savepoint; the in-process
+ runner commits the session).
+ """
+ from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
+ from airflow.models.taskinstance import TaskInstance
+
+ payload = row.payload
+ ti_key = payload["ti_key"]
+ ti = session.scalar(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == ti_key["dag_id"],
+ TaskInstance.run_id == ti_key["run_id"],
+ TaskInstance.task_id == ti_key["task_id"],
+ TaskInstance.map_index == ti_key["map_index"],
+ )
+ )
+ if ti is not None:
+ task_outlets = [AssetProfile.model_validate(outlet) for outlet in
payload["task_outlets"]]
+ TaskInstance.register_asset_changes_in_db(ti, task_outlets,
payload["outlet_events"], session=session)
Review Comment:
We should log (info level should be enough) that if the TI is not found in
this case.
--
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]
