hkc-8010 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3700076454
##########
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:
Done. I renamed `_register_queued_asset_event` to public
`register_queued_asset_event` and updated the scheduler drain to call it
directly.
##########
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:
Done. I renamed `register_pending_asset_events` to
`register_queued_asset_events_for_dag_test`, updated the docstring to call out
that it is only for `dag.test`, and updated the sole call site.
##########
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:
Done. I removed the extra `ti_key` comment here since
`_enqueue_asset_events` already explains the architecture and the payload key
reads clearly on its own.
##########
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:
Done. If the live task instance cannot be resolved by natural key now, we
log it at `info` before dropping the queue row. I also updated the scheduler
test to assert that path.
--
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]
