fix(APDR): add asset events to partitioned DagRun (#61433)

* test: simplify test_partitioned_dag_run_with_customized_mapper

* test(scheduler_job): improve APDR test cases to include consumed_asset_event validation

* fix(APDR): add asset events to partitioned DagRun

* fixup! fix(APDR): add asset events to partitioned DagRun

* fixup! test: simplify test_partitioned_dag_run_with_customized_mapper

* fixup! test(scheduler_job): improve APDR test cases to include consumed_asset_event validation

* fixup! test(scheduler_job): improve APDR test cases to include consumed_asset_event validation

* fixup! test(scheduler_job): improve APDR test cases to include consumed_asset_event validation

Author: Lee-W

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -1848,6 +1848,13 @@ 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,
+                )
+            )
+            dag_run.consumed_asset_events.extend(asset_events)
             session.flush()
             apdr.created_dag_run_id = dag_run.id
             session.flush()

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -8810,46 +8810,15 @@ def test_partitioned_dag_run_with_customized_mapper(
     runner = SchedulerJobRunner(
         job=Job(job_type=SchedulerJobRunner.job_type), executors=[MockExecutor(do_update=False)]
     )
-
-    with dag_maker(dag_id="asset-event-producer", schedule=None, session=session) as dag:
-        EmptyOperator(task_id="hi", outlets=[asset_1])
-
-    dr = dag_maker.create_dagrun(partition_key="this-is-not-key-1-before-mapped", session=session)
-    [ti] = dr.get_task_instances(session=session)
-    session.commit()
-
-    serialized_outlets = dag.get_task("hi").outlets
     with custom_partition_mapper_patch():
-        TaskInstance.register_asset_changes_in_db(
-            ti=ti,
-            task_outlets=[o.asprofile() for o in serialized_outlets],
-            outlet_events=[],
+        apdr = _produce_and_register_asset_event(
+            dag_id="asset-event-producer",
+            asset=asset_1,
+            partition_key="this-is-not-key-1-before-mapped",
             session=session,
+            dag_maker=dag_maker,
+            expected_partition_key="key-1",
         )
-    session.commit()
-
-    event = session.scalar(
-        select(AssetEvent).where(
-            AssetEvent.source_dag_id == dag.dag_id,
-            AssetEvent.source_run_id == dr.run_id,
-        )
-    )
-    assert event is not None
-    assert event.partition_key == "this-is-not-key-1-before-mapped"
-
-    apdr = session.scalar(
-        select(AssetPartitionDagRun)
-        .join(
-            PartitionedAssetKeyLog,
-            PartitionedAssetKeyLog.asset_partition_dag_run_id == AssetPartitionDagRun.id,
-        )
-        .where(PartitionedAssetKeyLog.asset_event_id == event.id)
-    )
-    assert apdr is not None
-    assert apdr.created_dag_run_id is None
-    assert apdr.partition_key == "key-1"
-
-    with custom_partition_mapper_patch():
         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,
@@ -8858,6 +8827,13 @@ def test_partitioned_dag_run_with_customized_mapper(
     assert len(partition_dags) == 1
     assert partition_dags == {"asset-event-consumer"}
 
+    dag_run = session.scalar(select(DagRun).where(DagRun.id == apdr.created_dag_run_id))
+    assert dag_run is not None
+    asset_event = dag_run.consumed_asset_events[0]
+    assert asset_event.source_task_id == "hi"
+    assert asset_event.source_dag_id == "asset-event-producer"
+    assert asset_event.source_run_id == "test"
+
 
 @pytest.mark.need_serialized_dag
 @pytest.mark.usefixtures("clear_asset_partition_rows")
@@ -8931,6 +8907,13 @@ def test_consumer_dag_listen_to_two_partitioned_asset(
     assert len(partition_dags) == 1
     assert partition_dags == {"asset-event-consumer"}
 
+    dag_run = session.scalar(select(DagRun).where(DagRun.id == apdr.created_dag_run_id))
+    assert dag_run is not None
+    for asset_event in dag_run.consumed_asset_events:
+        assert asset_event.source_task_id == "hi"
+        assert "asset-event-producer-" in asset_event.source_dag_id
+        assert asset_event.source_run_id == "test"
+
 
 @pytest.mark.need_serialized_dag
 @pytest.mark.usefixtures("clear_asset_partition_rows")
@@ -8995,3 +8978,10 @@ def test_consumer_dag_listen_to_two_partitioned_asset_with_key_1_mapper(
     assert apdr.created_dag_run_id is not None
     assert len(partition_dags) == 1
     assert partition_dags == {"asset-event-consumer"}
+
+    dag_run = session.scalar(select(DagRun).where(DagRun.id == apdr.created_dag_run_id))
+    assert dag_run is not None
+    for asset_event in dag_run.consumed_asset_events:
+        assert asset_event.source_task_id == "hi"
+        assert "asset-event-producer-" in asset_event.source_dag_id
+        assert asset_event.source_run_id == "test"