This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4acd1a9a1a2 Reject mismatched rollup mapper and window pairings at Dag
parse time (#69456)
4acd1a9a1a2 is described below
commit 4acd1a9a1a2001a0c83c8d3f0cd308f2b5d9dfd3
Author: Wei Lee <[email protected]>
AuthorDate: Tue Jul 7 12:22:28 2026 +0800
Reject mismatched rollup mapper and window pairings at Dag parse time
(#69456)
---
airflow-core/src/airflow/partition_mappers/base.py | 17 +++---
.../src/airflow/partition_mappers/temporal.py | 2 +
.../tests/unit/partition_mappers/test_base.py | 66 +++++++++++++++++++++-
.../sdk/definitions/partition_mappers/base.py | 9 ++-
.../task_sdk/definitions/test_partition_mappers.py | 53 +++++++++++++----
5 files changed, 122 insertions(+), 25 deletions(-)
diff --git a/airflow-core/src/airflow/partition_mappers/base.py
b/airflow-core/src/airflow/partition_mappers/base.py
index b0edd02f645..c773ef7e867 100644
--- a/airflow-core/src/airflow/partition_mappers/base.py
+++ b/airflow-core/src/airflow/partition_mappers/base.py
@@ -38,10 +38,12 @@ class PartitionMapper(ABC):
is_rollup: ClassVar[bool] = False
- # Class-level default so the attribute resolves even on subclasses whose
- # __init__ (e.g. attrs-generated in plugins) never calls this base
__init__.
max_downstream_keys: int | None = None
+ # Declared result type of decode_downstream; RollupMapper rejects windows
+ # whose expected_decoded_type does not match. Temporal mappers use
datetime.
+ expected_decoded_type: ClassVar[type] = str
+
def __init__(self, *, max_downstream_keys: int | None = None) -> None:
if max_downstream_keys is not None and (
not isinstance(max_downstream_keys, int) or max_downstream_keys < 1
@@ -168,16 +170,13 @@ class RollupMapper(PartitionMapper):
wait_policy: WaitPolicy | None = None,
max_downstream_keys: int | None = None,
) -> None:
- decode_overridden = type(upstream_mapper).decode_downstream is not
PartitionMapper.decode_downstream
- if not decode_overridden and window.expected_decoded_type is not str:
+ if upstream_mapper.expected_decoded_type is not
window.expected_decoded_type:
raise TypeError(
f"{type(window).__name__} expects decoded values of type "
f"{window.expected_decoded_type.__name__!r}, but "
- f"{type(upstream_mapper).__name__} does not override "
- f"'decode_downstream' (base default returns str). Pair the
window "
- f"with an upstream mapper that decodes to "
- f"{window.expected_decoded_type.__name__}, or use a window
whose "
- f"'expected_decoded_type' accepts str."
+ f"{type(upstream_mapper).__name__} decodes to "
+ f"{upstream_mapper.expected_decoded_type.__name__!r}. Pair a
window and an "
+ f"upstream mapper whose 'expected_decoded_type' match."
)
if wait_policy is None:
wait_policy = WaitForAll()
diff --git a/airflow-core/src/airflow/partition_mappers/temporal.py
b/airflow-core/src/airflow/partition_mappers/temporal.py
index c16801c97b1..09f1897a590 100644
--- a/airflow-core/src/airflow/partition_mappers/temporal.py
+++ b/airflow-core/src/airflow/partition_mappers/temporal.py
@@ -155,6 +155,8 @@ def _compile_output_format_regex(
class _BaseTemporalMapper(PartitionMapper):
"""Base class for Temporal Partition Mappers."""
+ expected_decoded_type: ClassVar[type] = datetime
+
default_output_format: str
def __init__(
diff --git a/airflow-core/tests/unit/partition_mappers/test_base.py
b/airflow-core/tests/unit/partition_mappers/test_base.py
index 5e730e49dd8..e1aef80909b 100644
--- a/airflow-core/tests/unit/partition_mappers/test_base.py
+++ b/airflow-core/tests/unit/partition_mappers/test_base.py
@@ -22,9 +22,10 @@ from datetime import datetime, timezone
import pytest
from airflow.partition_mappers.base import PartitionMapper, RollupMapper
+from airflow.partition_mappers.fixed_key import FixedKeyMapper
from airflow.partition_mappers.identity import IdentityMapper
from airflow.partition_mappers.temporal import StartOfDayMapper
-from airflow.partition_mappers.window import DayWindow
+from airflow.partition_mappers.window import DayWindow, SegmentWindow
from airflow.serialization.decoders import decode_partition_mapper
from airflow.serialization.encoders import encode_partition_mapper
from airflow.serialization.enums import Encoding
@@ -146,6 +147,69 @@ class TestRollupMapperInit:
# Should not raise.
RollupMapper(upstream_mapper=_StringOnlyMapper(),
window=_AlphaWindow())
+ def test_accepts_decode_overriding_str_mapper_with_str_window(self):
+ """Overriding decode/encode while still decoding to str is a legal
str-window pairing.
+
+ The guard must key off the declared ``expected_decoded_type`` (str by
+ default), not off whether ``decode_downstream`` is overridden.
+ """
+
+ class _CasefoldMapper(PartitionMapper):
+ def to_downstream(self, key):
+ return key.casefold()
+
+ def decode_downstream(self, downstream_key):
+ return downstream_key.casefold()
+
+ def encode_upstream(self, decoded_upstream):
+ return str(decoded_upstream)
+
+ # Should not raise.
+ RollupMapper(upstream_mapper=_CasefoldMapper(),
window=SegmentWindow(["us", "eu"]))
+
+ @pytest.mark.parametrize(
+ ("upstream_mapper_factory", "window_factory", "match"),
+ [
+ pytest.param(
+ lambda: FixedKeyMapper("all"),
+ lambda: SegmentWindow(["us", "eu"]),
+ None,
+ id="str_mapper-str_window-valid",
+ ),
+ pytest.param(
+ StartOfDayMapper,
+ DayWindow,
+ None,
+ id="datetime_mapper-datetime_window-valid",
+ ),
+ pytest.param(
+ lambda: FixedKeyMapper("all"),
+ DayWindow,
+ "DayWindow expects decoded values of type 'datetime'",
+ id="str_mapper-datetime_window-invalid",
+ ),
+ pytest.param(
+ StartOfDayMapper,
+ lambda: SegmentWindow(["us"]),
+ "SegmentWindow expects decoded values of type 'str'",
+ id="datetime_mapper-str_window-invalid",
+ ),
+ ],
+ )
+ def test_upstream_mapper_window_type_pairing(self,
upstream_mapper_factory, window_factory, match):
+ """RollupMapper's guard must catch a decoded-type mismatch in either
direction.
+
+ Guards the ``datetime_mapper-str_window`` direction in particular: left
+ unchecked, that pairing survives construction and serialization, then
+ raises ``AttributeError`` inside ``to_upstream`` at scheduler tick
time.
+ """
+ if match is None:
+ # Should not raise.
+ RollupMapper(upstream_mapper=upstream_mapper_factory(),
window=window_factory())
+ else:
+ with pytest.raises(TypeError, match=match):
+ RollupMapper(upstream_mapper=upstream_mapper_factory(),
window=window_factory())
+
class TestPartitionMapperMaxDownstreamKeysValidator:
"""Verify the max_downstream_keys validator on the PartitionMapper base
class.
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 079b85fc021..1336f29d54d 100644
--- a/task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py
+++ b/task-sdk/src/airflow/sdk/definitions/partition_mappers/base.py
@@ -82,12 +82,11 @@ class RollupMapper(PartitionMapper):
# would otherwise be swallowed by the bare ``except`` in
# ``_create_dagruns_for_partitioned_asset_dags`` and surface only as
# "Failed to deserialize Dag" spam).
- if self.upstream_mapper.expected_decoded_type is str and
self.window.expected_decoded_type is not str:
+ if self.upstream_mapper.expected_decoded_type is not
self.window.expected_decoded_type:
raise TypeError(
f"{type(self.window).__name__} expects decoded values of type "
f"{self.window.expected_decoded_type.__name__!r}, but "
- f"{type(self.upstream_mapper).__name__} decodes to 'str' (SDK
PartitionMapper default). "
- f"Pair the window with an upstream mapper whose
'expected_decoded_type' is "
- f"{self.window.expected_decoded_type.__name__}, or use a
window whose "
- f"'expected_decoded_type' accepts str."
+ f"{type(self.upstream_mapper).__name__} decodes to "
+ f"{self.upstream_mapper.expected_decoded_type.__name__!r}.
Pair a window and an "
+ f"upstream mapper whose 'expected_decoded_type' match."
)
diff --git a/task-sdk/tests/task_sdk/definitions/test_partition_mappers.py
b/task-sdk/tests/task_sdk/definitions/test_partition_mappers.py
index ca811115cb7..e588dc618d3 100644
--- a/task-sdk/tests/task_sdk/definitions/test_partition_mappers.py
+++ b/task-sdk/tests/task_sdk/definitions/test_partition_mappers.py
@@ -192,14 +192,47 @@ class TestSdkSegmentWindow:
class TestSdkCategoricalRollupGuard:
- """SDK-side RollupMapper guard mirrors core: str mapper + str window
passes."""
+ """SDK-side RollupMapper guard mirrors core: the decoded type must match
in either direction."""
- def test_fixed_key_with_segment_window_does_not_raise(self):
- # SDK guard: FixedKeyMapper.expected_decoded_type is str,
- # SegmentWindow.expected_decoded_type is str -> guard passes.
- RollupMapper(upstream_mapper=FixedKeyMapper("all"),
window=SegmentWindow(["us", "eu"]))
-
- def test_str_mapper_with_datetime_window_raises(self):
- # SDK guard: FixedKeyMapper (str) + DayWindow (datetime) -> raise.
- with pytest.raises(TypeError, match="DayWindow expects decoded values
of type 'datetime'"):
- RollupMapper(upstream_mapper=FixedKeyMapper("all"),
window=DayWindow())
+ @pytest.mark.parametrize(
+ ("upstream_mapper_factory", "window_factory", "match"),
+ [
+ pytest.param(
+ lambda: FixedKeyMapper("all"),
+ lambda: SegmentWindow(["us", "eu"]),
+ None,
+ id="str_mapper-str_window-valid",
+ ),
+ pytest.param(
+ StartOfDayMapper,
+ DayWindow,
+ None,
+ id="datetime_mapper-datetime_window-valid",
+ ),
+ pytest.param(
+ lambda: FixedKeyMapper("all"),
+ DayWindow,
+ "DayWindow expects decoded values of type 'datetime'",
+ id="str_mapper-datetime_window-invalid",
+ ),
+ pytest.param(
+ StartOfDayMapper,
+ lambda: SegmentWindow(["us"]),
+ "SegmentWindow expects decoded values of type 'str'",
+ id="datetime_mapper-str_window-invalid",
+ ),
+ ],
+ )
+ def test_upstream_mapper_window_type_pairing(self,
upstream_mapper_factory, window_factory, match):
+ """RollupMapper's guard must catch a decoded-type mismatch in either
direction.
+
+ Guards the ``datetime_mapper-str_window`` direction in particular: left
+ unchecked, that pairing survives construction and serialization, then
+ raises ``AttributeError`` inside ``to_upstream`` at scheduler tick
time.
+ """
+ if match is None:
+ # Should not raise.
+ RollupMapper(upstream_mapper=upstream_mapper_factory(),
window=window_factory())
+ else:
+ with pytest.raises(TypeError, match=match):
+ RollupMapper(upstream_mapper=upstream_mapper_factory(),
window=window_factory())
