diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 3f2c7c8a7736d..c0a57a8b9c9dd 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, @@ -138,6 +140,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.elements import ColumnElement from sqlalchemy.sql.selectable import Subquery @@ -163,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. @@ -2425,13 +2445,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() @@ -2688,26 +2710,33 @@ def _create_dag_runs_asset_triggered( ) ), ) - asset_events = 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()) + selected_asset_events = ( + select( + AssetEvent.id.label("event_id"), + AssetEvent.timestamp.label("timestamp"), ) + .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() + ), + ) + .cte() ) - if asset_events: - triggered_date = timezone.coerce_datetime(max(event.timestamp for event in asset_events)) + triggered_date, max_selected_event_id = session.execute( + select( + func.max(selected_asset_events.c.timestamp), + func.max(selected_asset_events.c.event_id), + ) + ).one() + if triggered_date is not None and max_selected_event_id is not None: + triggered_date = timezone.coerce_datetime(triggered_date) self.log.debug( "Creating asset-triggered DagRun for '%s': %d queued assets, triggered_date=%s", dag.dag_id, @@ -2733,12 +2762,32 @@ 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=select(selected_asset_events.c.event_id) + .where( + selected_asset_events.c.timestamp <= triggered_date, + selected_asset_events.c.event_id <= max_selected_event_id, + ) + .order_by( + selected_asset_events.c.timestamp.asc(), + selected_asset_events.c.event_id.asc(), + ), + 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, ) else: self.log.info( diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 440e6e2648c46..26759ade47128 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -38,7 +38,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 @@ -64,7 +64,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, @@ -196,6 +196,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 @@ -5658,6 +5681,31 @@ 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) + assert dag_run.consumed_asset_events == [] + + 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): """ @@ -5731,7 +5779,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.""" @@ -5939,6 +5988,84 @@ def test_create_dag_runs_asset_triggered_skips_stale_triggered_date(self, sessio # We do not create a new DagRun since the ADRQ has already been consumed assert session.scalars(select(DagRun).where(DagRun.dag_id == dag_model.dag_id)).one_or_none() is None + @pytest.mark.need_serialized_dag + @mock.patch("airflow.jobs.scheduler_job_runner._associate_asset_events_with_dag_run", autospec=True) + def test_asset_event_queued_during_association_waits_for_next_run( + self, mock_associate_asset_events, session, dag_maker + ): + asset = Asset(name="event-queued-during-association") + with dag_maker( + dag_id="event-queued-during-association-consumer", + schedule=[asset], + catchup=True, + session=session, + ): + pass + dag_model = dag_maker.dag_model + asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri)) + first_timestamp = timezone.utcnow() + first_event = AssetEvent(asset_id=asset_id, timestamp=first_timestamp) + session.add(first_event) + session.flush() + session.add( + AssetDagRunQueue( + target_dag_id=dag_model.dag_id, + asset_id=asset_id, + asset_event_id=first_event.id, + ) + ) + session.flush() + + second_timestamp = first_timestamp + timedelta(minutes=1) + second_event_ids = [] + + def queue_second_event_then_associate(*, dag_run, asset_event_ids_select, session): + if not second_event_ids: + second_event = AssetEvent(asset_id=asset_id, timestamp=second_timestamp) + session.add(second_event) + session.flush() + second_event_ids.append(second_event.id) + session.add( + AssetDagRunQueue( + target_dag_id=dag_model.dag_id, + asset_id=asset_id, + asset_event_id=second_event.id, + ) + ) + session.flush() + _associate_asset_events_with_dag_run( + dag_run=dag_run, + asset_event_ids_select=asset_event_ids_select, + session=session, + ) + + mock_associate_asset_events.side_effect = queue_second_event_then_associate + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + + self.job_runner._create_dag_runs_asset_triggered(dag_models=[dag_model], session=session) + + first_run = session.scalars(select(DagRun).where(DagRun.dag_id == dag_model.dag_id)).one() + assert [event.id for event in first_run.consumed_asset_events] == [first_event.id] + assert first_run.run_after == first_timestamp + assert ( + session.scalars( + select(AssetDagRunQueue.asset_event_id).where( + AssetDagRunQueue.target_dag_id == dag_model.dag_id + ) + ).all() + == second_event_ids + ) + + self.job_runner._create_dag_runs_asset_triggered(dag_models=[dag_model], session=session) + + runs = session.scalars( + select(DagRun).where(DagRun.dag_id == dag_model.dag_id).order_by(DagRun.run_after) + ).all() + assert len(runs) == 2 + assert [event.id for event in runs[1].consumed_asset_events] == second_event_ids + assert runs[1].run_after == second_timestamp + @pytest.mark.need_serialized_dag def test_create_dag_runs_asset_triggered_deletes_only_selected_adrq_rows( self, session: Session, dag_maker @@ -11589,7 +11716,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