Validate partition mapper source keys before scheduling

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

@@ -589,19 +589,21 @@ def _queue_partitioned_dags(
 
             mapper = timetable.get_partition_mapper(name=asset_model.name, uri=asset_model.uri)
             try:
-                # We'll need to catch every possible exception happen when mapping partition_key.
+                # We'll need to catch every possible exception that happens
+                # when validating or mapping partition_key.
+                mapper.validate_source_key(partition_key)
                 target_key = mapper.to_downstream(partition_key)
             except Exception as err:
                 log.exception(
-                    "Could not map partition key for asset in target Dag. "
+                    "Could not validate or map partition key for asset in target Dag. "
                     "This likely indicates the target Dag's partition mapper "
                     "is misconfigured, or does not support this partition key.",
                     partition_key=partition_key,
                     asset=asset_model,
                     target_dag=target_dag,
                 )
                 log_extra = (
-                    f"Could not map partition_key '{partition_key}' for asset "
+                    f"Could not validate or 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 "

airflow-core/src/airflow/partition_mappers/allowed_key.py

@@ -29,11 +29,17 @@ def __init__(self, allowed_keys: list[str], *, max_downstream_keys: int | None =
         super().__init__(max_downstream_keys=max_downstream_keys)
         self.allowed_keys = allowed_keys
 
-    def to_downstream(self, key: str) -> str:
+    def _validate_key(self, key: str) -> None:
         if key not in self.allowed_keys:
             raise ValueError(f"Key {key!r} not in allowed keys {self.allowed_keys}")
+
+    def to_downstream(self, key: str) -> str:
+        self._validate_key(key)
         return key
 
+    def validate_source_key(self, key: str) -> None:
+        self._validate_key(key)
+
     def serialize(self) -> dict[str, Any]:
         data: dict[str, Any] = {"allowed_keys": self.allowed_keys}
         if self.max_downstream_keys is not None:

airflow-core/src/airflow/partition_mappers/base.py

@@ -138,6 +138,16 @@ def carry_partition_date(self, source_partition_date: datetime | None) -> dateti
         """
         return None
 
+    def validate_source_key(self, key: str) -> None:
+        """
+        Validate that a source partition key is in a supported format.
+
+        The default implementation accepts any key. Mappers that parse or
+        otherwise expect structured source keys should override this to reject
+        malformed-but-parseable inputs before they are normalized.
+        """
+        return None
+
     def serialize(self) -> dict[str, Any]:
         if self.max_downstream_keys is None:
             return {}

airflow-core/src/airflow/partition_mappers/chain.py

@@ -40,35 +40,52 @@ def __init__(
         super().__init__(max_downstream_keys=max_downstream_keys)
         self.mappers = [mapper0, mapper1, *mappers]
 
+    def _map_once(
+        self,
+        mapper: PartitionMapper,
+        keys: list[str],
+        *,
+        validate_source: bool = False,
+    ) -> list[str]:
+        next_keys: list[str] = []
+        for current_key in keys:
+            if validate_source:
+                mapper.validate_source_key(current_key)
+
+            mapped = mapper.to_downstream(current_key)
+            if not isinstance(mapped, (str, Iterable)):
+                raise TypeError(
+                    f"ChainMapper child mappers must return a string or iterable of strings, "
+                    f"but {type(mapper).__name__} returned {type(mapped).__name__}"
+                )
+
+            if isinstance(mapped, str):
+                next_keys.append(mapped)
+            elif isinstance(mapped, Iterable):
+                for mapped_key in mapped:
+                    if not isinstance(mapped_key, str):
+                        raise TypeError(
+                            f"ChainMapper child mappers must return an iterable of strings, "
+                            f"but {type(mapper).__name__} yielded {type(mapped_key).__name__}"
+                        )
+                    next_keys.append(mapped_key)
+        return next_keys
+
     def to_downstream(self, key: str) -> str | Iterable[str]:
         keys: list[str] = [key]
         for mapper in self.mappers:
-            next_keys: list[str] = []
-            for current_key in keys:
-                mapped = mapper.to_downstream(current_key)
-                if not isinstance(mapped, (str, Iterable)):
-                    raise TypeError(
-                        f"ChainMapper child mappers must return a string or iterable of strings, "
-                        f"but {type(mapper).__name__} returned {type(mapped).__name__}"
-                    )
-
-                if isinstance(mapped, str):
-                    next_keys.append(mapped)
-                elif isinstance(mapped, Iterable):
-                    for mapped_key in mapped:
-                        if not isinstance(mapped_key, str):
-                            raise TypeError(
-                                f"ChainMapper child mappers must return an iterable of strings, "
-                                f"but {type(mapper).__name__} yielded {type(mapped_key).__name__}"
-                            )
-                        next_keys.append(mapped_key)
-            keys = next_keys
+            keys = self._map_once(mapper, keys)
         return keys[0] if len(keys) == 1 else keys
 
     def to_partition_date(self, downstream_key: str) -> datetime | None:
         # The last mapper in the chain formats the final downstream key, so it owns the anchor.
         return self.mappers[-1].to_partition_date(downstream_key)
 
+    def validate_source_key(self, key: str) -> None:
+        keys: list[str] = [key]
+        for mapper in self.mappers:
+            keys = self._map_once(mapper, keys, validate_source=True)
+
     def serialize(self) -> dict[str, Any]:
         from airflow.serialization.encoders import encode_partition_mapper
 

airflow-core/src/airflow/partition_mappers/product.py

@@ -41,10 +41,14 @@ def __init__(
         self.mappers = [mapper0, mapper1, *mappers]
         self.delimiter = delimiter
 
-    def to_downstream(self, key: str) -> str:
+    def _split_segments(self, key: str) -> list[str]:
         segments = key.split(self.delimiter)
         if len(segments) != len(self.mappers):
             raise ValueError(f"Expected {len(self.mappers)} segments in key, got {len(segments)}")
+        return segments
+
+    def to_downstream(self, key: str) -> str:
+        segments = self._split_segments(key)
         results: list[str] = []
         for mapper, segment in zip(self.mappers, segments):
             result = mapper.to_downstream(segment)
@@ -82,6 +86,11 @@ def to_partition_date(self, downstream_key: str) -> datetime | None:
             return anchors[0]
         return None
 
+    def validate_source_key(self, key: str) -> None:
+        segments = self._split_segments(key)
+        for mapper, segment in zip(self.mappers, segments):
+            mapper.validate_source_key(segment)
+
     def serialize(self) -> dict[str, Any]:
         from airflow.serialization.encoders import encode_partition_mapper
 

airflow-core/src/airflow/partition_mappers/temporal.py

@@ -174,15 +174,26 @@ def __init__(
             timezone = parse_timezone(timezone)
         self._timezone = timezone
 
-    def to_downstream(self, key: str) -> str:
-        dt = datetime.strptime(key, self.input_format)
+    def _parse_source_key(self, key: str) -> datetime:
+        return datetime.strptime(key, self.input_format)
+
+    def _normalize_timezone(self, dt: datetime) -> datetime:
         if dt.tzinfo is None:
-            dt = make_aware(dt, self._timezone)
-        else:
-            dt = dt.astimezone(self._timezone)
+            return make_aware(dt, self._timezone)
+        return dt.astimezone(self._timezone)
+
+    def to_downstream(self, key: str) -> str:
+        dt = self._normalize_timezone(self._parse_source_key(key))
         normalized = self.normalize(dt)
         return self.format(normalized)
 
+    def validate_source_key(self, key: str) -> None:
+        canonical = self._parse_source_key(key).strftime(self.input_format)
+        if canonical != key:
+            raise ValueError(
+                f"Partition key {key!r} does not round-trip with input format {self.input_format!r}"
+            )
+
     @abstractmethod
     def normalize(self, dt: datetime) -> datetime:
         """Return canonical start datetime for the partition."""

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

@@ -10658,14 +10658,89 @@ def test_partitioned_dag_run_with_invalid_mapping(dag_maker: DagMaker, session:
     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' "
+        "Could not validate or 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_rejects_non_canonical_partition_key(dag_maker: DagMaker, session: Session):
+    session.execute(delete(Log))
+    asset_1 = Asset(name="asset-1")
+    with dag_maker(
+        dag_id="asset-event-consumer",
+        schedule=PartitionedAssetTimetable(
+            assets=asset_1,
+            default_partition_mapper=StartOfHourMapper(),
+        ),
+        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=CronPartitionTimetable("* * * * *", timezone="UTC"),
+        session=session,
+    ) as dag:
+        EmptyOperator(task_id="hi", outlets=[asset_1])
+
+    partition_key = "2026-2-10T4:30:45"
+    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 validate or map partition_key '2026-2-10T4:30:45' "
+        "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: Partition key '2026-2-10T4:30:45' does not round-trip with input format "
+        "'%Y-%m-%dT%H:%M:%S'"
+    )
+
+
 @pytest.mark.db_test
 def test_create_dag_runs_partitioned_timetable_skips_when_next_fields_none(session):
     """

airflow-core/tests/unit/partition_mappers/test_allowed_key.py

@@ -35,6 +35,15 @@ def test_to_downstream_invalid_key(self):
         with pytest.raises(ValueError, match="not in allowed keys"):
             pm.to_downstream("apac")
 
+    def test_validate_source_key(self):
+        pm = AllowedKeyMapper(["us", "eu", "apac"])
+        pm.validate_source_key("apac")
+
+    def test_validate_source_key_invalid_key(self):
+        pm = AllowedKeyMapper(["us", "eu"])
+        with pytest.raises(ValueError, match="not in allowed keys"):
+            pm.validate_source_key("apac")
+
     def test_serialize(self):
         pm = AllowedKeyMapper(["a", "b", "c"])
         assert pm.serialize() == {"allowed_keys": ["a", "b", "c"]}

airflow-core/tests/unit/partition_mappers/test_chain.py

@@ -61,6 +61,15 @@ def test_to_downstream(self):
     def test_to_partition_date_delegates_to_last_mapper(self, chain, downstream_key, expected):
         assert chain.to_partition_date(downstream_key) == expected
 
+    def test_validate_source_key(self):
+        sm = ChainMapper(StartOfHourMapper(), StartOfDayMapper(input_format="%Y-%m-%dT%H"))
+        sm.validate_source_key("2024-01-15T10:30:00")
+
+    def test_validate_source_key_rejects_non_canonical_input(self):
+        sm = ChainMapper(StartOfHourMapper(), StartOfDayMapper(input_format="%Y-%m-%dT%H"))
+        with pytest.raises(ValueError, match="does not round-trip"):
+            sm.validate_source_key("2024-1-15T10:30:00")
+
     def test_to_downstream_invalid_non_iterable_return(self):
         sm = ChainMapper(IdentityMapper(), _InvalidReturnMapper())
         with pytest.raises(TypeError, match="must return a string or iterable of strings"):

airflow-core/tests/unit/partition_mappers/test_product.py

@@ -44,6 +44,15 @@ def test_to_downstream_single_segment_for_two_mappers(self):
         with pytest.raises(ValueError, match="Expected 2 segments"):
             pm.to_downstream("2024-01-15T10:30:00")
 
+    def test_validate_source_key(self):
+        pm = ProductMapper(StartOfHourMapper(), StartOfDayMapper())
+        pm.validate_source_key("2024-01-15T10:30:00|2024-01-15T10:30:00")
+
+    def test_validate_source_key_rejects_non_canonical_segment(self):
+        pm = ProductMapper(StartOfHourMapper(), StartOfDayMapper())
+        with pytest.raises(ValueError, match="does not round-trip"):
+            pm.validate_source_key("2024-1-15T10:30:00|2024-01-15T10:30:00")
+
     def test_custom_delimiter(self):
         pm = ProductMapper(StartOfHourMapper(), StartOfDayMapper(), delimiter="::")
         assert pm.to_downstream("2024-01-15T10:30:00::2024-01-15T10:30:00") == "2024-01-15T10::2024-01-15"

airflow-core/tests/unit/partition_mappers/test_temporal.py

@@ -215,6 +215,13 @@ def test_start_of_quarter_decode_rejects_trailing_garbage(self):
         with pytest.raises(ValueError, match="could not parse"):
             mapper.decode_downstream("2024-Q1 trailing-garbage")
 
+    def test_validate_source_key(self):
+        StartOfHourMapper().validate_source_key("2026-02-10T14:30:45")
+
+    def test_validate_source_key_rejects_non_canonical_input(self):
+        with pytest.raises(ValueError, match="does not round-trip"):
+            StartOfHourMapper().validate_source_key("2026-2-10T4:30:45")
+
 
 class TestSdkTemporalMappersTimezoneSerialization:
     """

task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py

@@ -50,6 +50,10 @@ class PartitionMapper:
         default=None, kw_only=True, validator=_validate_max_downstream_keys
     )
 
+    def validate_source_key(self, key: str) -> None:
+        """Validate that a source partition key is in a supported format."""
+        return None
+
 
 @attrs.define
 class RollupMapper(PartitionMapper):