diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 7d5f7f58020c7..d15ab59b0dec2 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -43,7 +43,9 @@ delete, exists, func, + insert, inspect, + literal, or_, select, text, @@ -85,6 +87,7 @@ PartitionedAssetKeyLog, TaskInletAssetReference, TaskOutletAssetReference, + association_table, ) from airflow.models.asset_state_store import AssetStateStoreModel from airflow.models.backfill import Backfill, BackfillDagRun @@ -138,6 +141,7 @@ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session from sqlalchemy.orm.interfaces import LoaderOption + from sqlalchemy.sql import Select from sqlalchemy.sql.selectable import Subquery from airflow._shared.logging.types import Logger @@ -162,6 +166,23 @@ MAX_PARTITION_DAG_RUNS_PER_LOOP = 500 +def _associate_asset_events_with_dag_run( + *, + dag_run: DagRun, + asset_event_ids_select: Select[tuple[int]], + session: Session, +) -> None: + """Associate selected asset events without materializing ORM objects.""" + selected_event_ids = asset_event_ids_select.subquery() + session.execute( + insert(association_table).from_select( + ["dag_run_id", "event_id"], + select(literal(dag_run.id), selected_event_ids.c.event_id), + ) + ) + session.expire(dag_run, ["consumed_asset_events"]) + + def _eager_load_dag_run_for_validation() -> tuple[LoaderOption, LoaderOption]: """ Eager-load DagRun relations required for execution API datamodel validation. @@ -2341,13 +2362,15 @@ def _create_dagruns_for_partitioned_asset_dags(self, session: Session) -> set[st creating_job_id=self.job.id, session=session, ) - asset_events = session.scalars( - select(AssetEvent).where( - PartitionedAssetKeyLog.asset_partition_dag_run_id == apdr.id, - PartitionedAssetKeyLog.asset_event_id == AssetEvent.id, - ) + asset_event_ids_select = select(AssetEvent.id.label("event_id")).where( + PartitionedAssetKeyLog.asset_partition_dag_run_id == apdr.id, + PartitionedAssetKeyLog.asset_event_id == AssetEvent.id, + ) + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids_select=asset_event_ids_select, + session=session, ) - dag_run.consumed_asset_events.extend(asset_events) session.flush() apdr.created_dag_run_id = dag_run.id session.flush() @@ -2619,27 +2642,25 @@ def _create_dag_runs_asset_triggered( ) event_window_floor.append(date.min) - asset_events = list( - session.scalars( - select(AssetEvent) - .where( - or_( - AssetEvent.asset_id.in_( - select(DagScheduleAssetReference.asset_id).where( - DagScheduleAssetReference.dag_id == dag.dag_id - ) - ), - AssetEvent.source_aliases.any( - AssetAliasModel.scheduled_dags.any( - DagScheduleAssetAliasReference.dag_id == dag.dag_id - ) + asset_event_ids_select = ( + select(AssetEvent.id.label("event_id")) + .where( + or_( + AssetEvent.asset_id.in_( + select(DagScheduleAssetReference.asset_id).where( + DagScheduleAssetReference.dag_id == dag.dag_id ), ), - AssetEvent.timestamp <= triggered_date, - AssetEvent.timestamp > func.coalesce(*event_window_floor), - ) - .order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc()) + AssetEvent.source_aliases.any( + AssetAliasModel.scheduled_dags.any( + DagScheduleAssetAliasReference.dag_id == dag.dag_id + ) + ), + ), + AssetEvent.timestamp <= triggered_date, + AssetEvent.timestamp > func.coalesce(*event_window_floor), ) + .order_by(AssetEvent.timestamp.asc(), AssetEvent.id.asc()) ) dag_run = dag.create_dagrun( @@ -2661,12 +2682,24 @@ def _create_dag_runs_asset_triggered( else None ) stats.incr("asset.triggered_dagruns", tags=prune_dict({"team_name": team_name})) - dag_run.consumed_asset_events.extend(asset_events) + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids_select=asset_event_ids_select, + session=session, + ) + consumed_asset_event_count = ( + session.scalar( + select(func.count()) + .select_from(association_table) + .where(association_table.c.dag_run_id == dag_run.id) + ) + or 0 + ) self.log.info( "Created asset-triggered DagRun for '%s': run_id=%s, consumed %d asset events", dag.dag_id, dag_run.run_id, - len(asset_events), + consumed_asset_event_count, ) # Delete only consumed ADRQ rows to avoid dropping newly queued events diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 1f3b6992c3e1c..79e67b98d8128 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -36,7 +36,7 @@ import psutil import pytest import time_machine -from sqlalchemy import delete, func, inspect, select, update +from sqlalchemy import delete, event as sqlalchemy_event, func, inspect, select, update from sqlalchemy.dialects import mysql from sqlalchemy.orm import joinedload @@ -57,7 +57,7 @@ from airflow.executors.executor_utils import ExecutorName from airflow.executors.local_executor import LocalExecutor from airflow.jobs.job import Job, run_job -from airflow.jobs.scheduler_job_runner import SchedulerJobRunner +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner, _associate_asset_events_with_dag_run from airflow.models.asset import ( AssetActive, AssetAliasModel, @@ -189,6 +189,29 @@ / "standard" / "example_dags" ) + + +@contextmanager +def assert_no_asset_event_materialization(session: Session) -> Generator[None, None, None]: + """Assert a scheduler operation does not select full AssetEvent entities.""" + materializing_queries = [] + + def track_orm_execute(orm_execute_state): + if orm_execute_state.is_select and any( + description.get("expr") is AssetEvent + for description in orm_execute_state.statement.column_descriptions + ): + materializing_queries.append(orm_execute_state.statement) + + sqlalchemy_event.listen(session, "do_orm_execute", track_orm_execute) + try: + yield + finally: + sqlalchemy_event.remove(session, "do_orm_execute", track_orm_execute) + + assert materializing_queries == [] + + DEFAULT_DATE = timezone.datetime(2016, 1, 1) DEFAULT_LOGICAL_DATE = timezone.coerce_datetime(DEFAULT_DATE) TRY_NUMBER = 1 @@ -5557,6 +5580,30 @@ def test_create_dag_runs(self, dag_maker): assert dr.start_date is None assert dr.creating_job_id == scheduler_job.id + def test_associate_asset_events_with_dag_run_keeps_relationship_coherent(self, session, dag_maker): + with dag_maker(dag_id="asset-event-association"): + pass + dag_run = dag_maker.create_dagrun() + asset = AssetModel(name="association-asset", uri="association-asset", group="asset") + session.add(asset) + session.flush() + asset_event = AssetEvent(asset_id=asset.id) + session.add(asset_event) + session.flush() + asset_event_id = asset_event.id + session.expunge(asset_event) + + with assert_queries_count(1, session=session): + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids_select=select(AssetEvent.id.label("event_id")).where( + AssetEvent.id == asset_event_id + ), + session=session, + ) + + assert [event.id for event in dag_run.consumed_asset_events] == [asset_event_id] + @pytest.mark.need_serialized_dag def test_create_dag_runs_assets(self, session, dag_maker): """ @@ -5632,7 +5679,8 @@ def test_create_dag_runs_assets(self, session, dag_maker): self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) with create_session() as session: - self.job_runner._create_dagruns_for_dags(session, session) + with assert_no_asset_event_materialization(session): + self.job_runner._create_dagruns_for_dags(session, session) def dict_from_obj(obj): """Get dict of column attrs from SqlAlchemy object.""" @@ -10727,7 +10775,8 @@ def test_partitioned_dag_run_with_customized_mapper( dag_maker=dag_maker, expected_partition_key="key-1", ) - partition_dags = runner._create_dagruns_for_partitioned_asset_dags(session=session) + with assert_no_asset_event_materialization(session): + partition_dags = runner._create_dagruns_for_partitioned_asset_dags(session=session) session.refresh(apdr) # Since asset event for Asset(name="asset-2") with key "key-1" has not yet been created, # no Dag run will be created