From e6d4911913f3943c75ab4d136556973f7f6383dd Mon Sep 17 00:00:00 2001 From: Sai Teja Desu Date: Mon, 18 May 2026 07:10:07 -0700 Subject: [PATCH] Validate partition mapper source keys before scheduling --- airflow-core/src/airflow/assets/manager.py | 8 +- .../airflow/partition_mappers/allowed_key.py | 8 +- .../src/airflow/partition_mappers/base.py | 10 +++ .../src/airflow/partition_mappers/chain.py | 57 +++++++++----- .../src/airflow/partition_mappers/product.py | 11 ++- .../src/airflow/partition_mappers/temporal.py | 21 +++-- .../tests/unit/jobs/test_scheduler_job.py | 77 ++++++++++++++++++- .../partition_mappers/test_allowed_key.py | 9 +++ .../unit/partition_mappers/test_chain.py | 9 +++ .../unit/partition_mappers/test_product.py | 9 +++ .../unit/partition_mappers/test_temporal.py | 7 ++ .../sdk/definitions/partition_mappers/base.py | 4 + 12 files changed, 199 insertions(+), 31 deletions(-) diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index c3491574e9148..dc324ae1e8eb0 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -584,11 +584,13 @@ 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, @@ -596,7 +598,7 @@ def _queue_partitioned_dags( 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 " diff --git a/airflow-core/src/airflow/partition_mappers/allowed_key.py b/airflow-core/src/airflow/partition_mappers/allowed_key.py index ff475eef0d29e..5eef3a8539a19 100644 --- a/airflow-core/src/airflow/partition_mappers/allowed_key.py +++ b/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: diff --git a/airflow-core/src/airflow/partition_mappers/base.py b/airflow-core/src/airflow/partition_mappers/base.py index c773ef7e86777..c553a8e2da3c6 100644 --- a/airflow-core/src/airflow/partition_mappers/base.py +++ b/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 {} diff --git a/airflow-core/src/airflow/partition_mappers/chain.py b/airflow-core/src/airflow/partition_mappers/chain.py index 6c0749feba221..7ade992d36fa9 100644 --- a/airflow-core/src/airflow/partition_mappers/chain.py +++ b/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 diff --git a/airflow-core/src/airflow/partition_mappers/product.py b/airflow-core/src/airflow/partition_mappers/product.py index 367ac4fda84d0..4924809c50a96 100644 --- a/airflow-core/src/airflow/partition_mappers/product.py +++ b/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 diff --git a/airflow-core/src/airflow/partition_mappers/temporal.py b/airflow-core/src/airflow/partition_mappers/temporal.py index 09f1897a590f0..66d0145f67340 100644 --- a/airflow-core/src/airflow/partition_mappers/temporal.py +++ b/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.""" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 59bc50a7f128c..b771e3fb8e34a 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -10508,7 +10508,7 @@ 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" @@ -10516,6 +10516,81 @@ def test_partitioned_dag_run_with_invalid_mapping(dag_maker: DagMaker, session: ) +@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): """ diff --git a/airflow-core/tests/unit/partition_mappers/test_allowed_key.py b/airflow-core/tests/unit/partition_mappers/test_allowed_key.py index 6fb47b5be31ec..468802e2dfeab 100644 --- a/airflow-core/tests/unit/partition_mappers/test_allowed_key.py +++ b/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"]} diff --git a/airflow-core/tests/unit/partition_mappers/test_chain.py b/airflow-core/tests/unit/partition_mappers/test_chain.py index a70230b223022..30691ec1f0e0f 100644 --- a/airflow-core/tests/unit/partition_mappers/test_chain.py +++ b/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"): diff --git a/airflow-core/tests/unit/partition_mappers/test_product.py b/airflow-core/tests/unit/partition_mappers/test_product.py index ce77441ef40e1..592764b05a248 100644 --- a/airflow-core/tests/unit/partition_mappers/test_product.py +++ b/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" diff --git a/airflow-core/tests/unit/partition_mappers/test_temporal.py b/airflow-core/tests/unit/partition_mappers/test_temporal.py index e777017b51a13..efa6d29ae52a1 100644 --- a/airflow-core/tests/unit/partition_mappers/test_temporal.py +++ b/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: """ diff --git a/task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py b/task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py index 1336f29d54dd9..d58269c1ba07e 100644 --- a/task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py +++ b/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):