fixup! feat(asset): add an Audit Log if the partition_key can not be mapped by target Dag partition_mapper

Author: Lee-W

airflow-core/src/airflow/assets/manager.py

@@ -342,8 +342,7 @@ def _queue_dagruns(
         task_instance: TaskInstance | None,
         session: Session,
     ) -> None:
-        log.debug("dags to queue", dags_to_queue=dags_to_queue)
-
+        log.debug("Dags to queue", dags_to_queue=dags_to_queue)
         if not dags_to_queue:
             return None
 
@@ -360,7 +359,6 @@ def _queue_dagruns(
         )
 
         non_partitioned_dags = dags_to_queue.difference(partition_dags)  # don't double process
-
         if not non_partitioned_dags:
             return None
 
@@ -422,20 +420,23 @@ def _queue_partitioned_dags(
                 raise RuntimeError(f"Could not find asset for asset_id={asset_id}")
 
             try:
+                # We'll need to catch every possible exception happen when mapping partition_key.
                 target_key = timetable.get_partition_mapper(
                     name=asset_model.name, uri=asset_model.uri
                 ).to_downstream(partition_key)
-            except Exception:
+            except Exception as err:
+                msg = (
+                    f"Could not map partition_key '{partition_key}' for asset "
+                    f"(name='{asset_model.name}', uri='{asset_model.uri}') in target Dag "
+                    f"'{target_dag.dag_id}'. This likely indicates that the partition "
+                    f"mapper in the target Dag is misconfigured or does not support this "
+                    f"partition key."
+                )
+                log.exception(msg)
                 session.add(
                     Log(
                         event="failed to map partition_key",
-                        extra=(
-                            f"Could not map partition_key '{partition_key}' for asset "
-                            f"(name='{asset_model.name}', uri='{asset_model.uri}') in target Dag "
-                            f"'{target_dag.dag_id}'. This likely indicates that the partition "
-                            f"mapper in the target Dag is misconfigured or does not support this "
-                            f"partition key."
-                        ),
+                        extra=(f"{msg}\n{type(err).__name__}: {err}"),
                         task_instance=task_instance,
                     )
                 )

airflow-core/tests/unit/assets/test_manager.py

@@ -44,6 +44,10 @@
 
 from unit.listeners import asset_listener
 
+if TYPE_CHECKING:
+    from tests_common.pytest_plugin import DagMaker
+
+
 pytestmark = pytest.mark.db_test
 
 
@@ -91,6 +95,38 @@ def test_register_asset_change_asset_doesnt_exist(self, mock_task_instance):
         mock_session.add.assert_not_called()
         mock_session.merge.assert_not_called()
 
+    @pytest.mark.need_serialized_dag
+    def test_register_asset_change_target_partition_mapper_raise_exception(
+        self,
+        session: Session,
+        dag_maker: DagMaker,
+    ):
+        bundle_name = "testing"
+        dag1 = DagModel(dag_id="dag1", is_stale=False, bundle_name=bundle_name)
+        dag2 = DagModel(
+            dag_id="dag2", is_stale=False, bundle_name=bundle_name, timetable_summary="Partitioned Asset"
+        )
+        session.add_all([dag1, dag2])
+
+        asm = AssetModel(uri="test://asset1/", name="test_asset_uri", group="asset")
+        session.add(asm)
+        asm.scheduled_dags = [DagScheduleAssetReference(dag_id=dag.dag_id) for dag in (dag1, dag2)]
+        session.execute(delete(AssetDagRunQueue))
+        session.flush()
+
+        asset_manager = AssetManager()
+        asset_manager.register_asset_change(
+            task_instance=mock_task_instance, asset=asm, session=session, partition_key="invalid"
+        )
+        session.flush()
+
+        # Ensure we've created an asset
+        assert (
+            session.scalar(select(func.count()).select_from(AssetEvent).where(AssetEvent.asset_id == asm.id))
+            == 1
+        )
+        assert session.scalar(select(func.count()).select_from(AssetDagRunQueue)) == 2
+
     @pytest.mark.usefixtures("dag_maker", "testing_dag_bundle")
     def test_register_asset_change(self, session, mock_task_instance):
         asset_manager = AssetManager()

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

@@ -84,7 +84,7 @@
 from airflow.providers.standard.operators.bash import BashOperator
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.providers.standard.triggers.file import FileDeleteTrigger
-from airflow.sdk import DAG, Asset, AssetAlias, AssetWatcher, IdentityMapper, task
+from airflow.sdk import DAG, Asset, AssetAlias, AssetWatcher, HourlyMapper, IdentityMapper, task
 from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
 from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
 from airflow.serialization.definitions.dag import SerializedDAG
@@ -8784,6 +8784,76 @@ def _produce_and_register_asset_event(
     return apdr
 
 
+@pytest.mark.need_serialized_dag
+@pytest.mark.usefixtures("clear_asset_partition_rows")
+def test_partitioned_dag_run_with_invalid_mapping(
+    dag_maker: DagMaker,
+    session: Session,
+):
+    asset_1 = Asset(name="asset-1")
+    with dag_maker(
+        dag_id="asset-event-consumer",
+        schedule=PartitionedAssetTimetable(
+            assets=asset_1,
+            default_partition_mapper=HourlyMapper(),
+        ),
+        session=session,
+    ):
+        EmptyOperator(task_id="hi")
+    session.commit()
+
+    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])
+
+    partition_key = "an invalid key for HourlyMapper"
+    dr = dag_maker.create_dagrun(partition_key=partition_key, session=session)
+    [ti] = dr.get_task_instances(session=session)
+    session.commit()
+
+    serialized_outlets = dag.get_task("hi").outlets
+    TaskInstance.register_asset_changes_in_db(
+        ti=ti,
+        task_outlets=[o.asprofile() for o in serialized_outlets],
+        outlet_events=[],
+        session=session,
+    )
+    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 == partition_key
+    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 None
+
+    partition_dags = runner._create_dagruns_for_partitioned_asset_dags(session=session)
+    assert len(partition_dags) == 0
+    assert partition_dags == set()
+
+    audit_log = session.scalar(select(Log))
+    assert audit_log is not None
+    assert audit_log.extra == (
+        "Could not map partition_key 'an invalid key for HourlyMapper' "
+        "for asset (name='asset-1', uri='asset-1') in target Dag 'asset-event-consumer'. "
+        "This likely indicates that the partition mapper in the target Dag is misconfigured or "
+        "does not support this partition key.\n"
+        "ValueError: time data 'an invalid key for HourlyMapper' does not match format '%Y-%m-%dT%H:%M:%S'"
+    )
+
+
 @pytest.mark.need_serialized_dag
 @pytest.mark.usefixtures("clear_asset_partition_rows")
 def test_partitioned_dag_run_with_customized_mapper(