You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by we...@apache.org on 2026/06/22 11:28:32 UTC
(airflow) branch main updated: Allow custom partition Window subclasses via plugin registry (#68717)
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 47dab8b8c09 Allow custom partition Window subclasses via plugin registry (#68717)
47dab8b8c09 is described below
commit 47dab8b8c0912e1ba173c73117e71959f518115f
Author: Wei Lee <we...@gmail.com>
AuthorDate: Mon Jun 22 19:28:11 2026 +0800
Allow custom partition Window subclasses via plugin registry (#68717)
---
.../example_dags/plugins/business_day_window.py | 64 ++++++++++++++++++++
airflow-core/src/airflow/plugins_manager.py | 9 +++
airflow-core/src/airflow/serialization/decoders.py | 13 ++--
airflow-core/src/airflow/serialization/encoders.py | 17 +++---
airflow-core/src/airflow/serialization/helpers.py | 18 ++++--
.../core_api/routes/public/test_plugins.py | 17 ++++--
.../tests/unit/partition_mappers/test_window.py | 69 ++++++++++++++++++----
.../tests/unit/plugins/test_plugins_manager.py | 29 ++++++++-
.../plugins_manager/plugins_manager.py | 3 +
9 files changed, 200 insertions(+), 39 deletions(-)
diff --git a/airflow-core/src/airflow/example_dags/plugins/business_day_window.py b/airflow-core/src/airflow/example_dags/plugins/business_day_window.py
new file mode 100644
index 00000000000..f3f984a5273
--- /dev/null
+++ b/airflow-core/src/airflow/example_dags/plugins/business_day_window.py
@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import calendar
+from collections.abc import Iterable
+from datetime import datetime, timedelta
+from typing import ClassVar
+
+from airflow.partition_mappers.window import Window
+from airflow.plugins_manager import AirflowPlugin
+
+
+# [START custom_window]
+class BusinessDayWindow(Window):
+ """
+ A calendar-month rollup window that yields only weekday (Mon–Fri) period-starts.
+
+ The built-in :class:`~airflow.partition_mappers.window.MonthWindow` yields every
+ calendar day in the month. ``BusinessDayWindow`` skips Saturdays and Sundays, so
+ a monthly downstream asset only waits for the business-day upstream partitions —
+ useful for financial or operational pipelines whose upstream data isn't produced
+ on weekends.
+
+ This class demonstrates registering a custom ``Window`` subclass via the
+ ``AirflowPlugin.windows`` registry: any plugin that lists it in ``windows = [...]``
+ makes it available to ``RollupMapper`` without modifying core Airflow.
+
+ *Assumes FORWARD direction and a day-1 month anchor* — the standard contract for
+ month-aligned temporal upstream mappers.
+ """
+
+ expected_decoded_type: ClassVar[type] = datetime
+
+ def to_upstream(self, period_start: datetime) -> Iterable[datetime]:
+ if period_start.day != 1:
+ raise ValueError(
+ f"BusinessDayWindow expects a period start on day 1 of the month, "
+ f"got {period_start.isoformat()}."
+ )
+ days_in_month = calendar.monthrange(period_start.year, period_start.month)[1]
+ return (day for i in range(days_in_month) if (day := period_start + timedelta(days=i)).weekday() < 5)
+
+
+class BusinessDayWindowPlugin(AirflowPlugin):
+ name = "business_day_window_plugin"
+ windows = [BusinessDayWindow]
+
+
+# [END custom_window]
diff --git a/airflow-core/src/airflow/plugins_manager.py b/airflow-core/src/airflow/plugins_manager.py
index c24b4bbd721..2aa93951e89 100644
--- a/airflow-core/src/airflow/plugins_manager.py
+++ b/airflow-core/src/airflow/plugins_manager.py
@@ -41,6 +41,7 @@ if TYPE_CHECKING:
from airflow.listeners.listener import ListenerManager
from airflow.models.deadline import DeadlineReferenceType
from airflow.partition_mappers.base import PartitionMapper
+ from airflow.partition_mappers.window import Window
from airflow.task.priority_strategy import PriorityWeightStrategy
from airflow.timetables.base import Timetable
@@ -287,6 +288,14 @@ def get_partition_mapper_plugins() -> dict[str, type[PartitionMapper]]:
}
+@cache
+def get_windows_plugins() -> dict[str, type[Window]]:
+ """Collect and get window classes registered by plugins."""
+ log.debug("Initialize extra window plugins")
+
+ return {qualname(window_cls): window_cls for plugin in _get_plugins()[0] for window_cls in plugin.windows}
+
+
@cache
def get_deadline_references_plugins() -> dict[str, type[DeadlineReferenceType]]:
"""Collect and get deadline reference classes registered by plugins."""
diff --git a/airflow-core/src/airflow/serialization/decoders.py b/airflow-core/src/airflow/serialization/decoders.py
index 12672818ed0..890a1296a4a 100644
--- a/airflow-core/src/airflow/serialization/decoders.py
+++ b/airflow-core/src/airflow/serialization/decoders.py
@@ -42,9 +42,9 @@ from airflow.serialization.definitions.deadline import (
from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
from airflow.serialization.helpers import (
WaitPolicyNotSupported,
- WindowNotSupported,
find_registered_custom_partition_mapper,
find_registered_custom_timetable,
+ find_registered_custom_window,
is_core_partition_mapper_import_path,
is_core_timetable_import_path,
is_core_wait_policy_import_path,
@@ -236,17 +236,18 @@ def decode_window(var: dict[str, Any]) -> Window:
"""
Decode a previously serialized :class:`Window`.
- Only built-in windows are accepted — a tampered serialized Dag naming a
- non-core import path is rejected up-front instead of being handed to
+ Custom windows must be registered via the ``windows`` plugin attribute;
+ unregistered import paths are rejected up-front instead of being handed to
``import_string``. See :func:`encode_window` for the matching encode-side
restriction.
:meta private:
"""
importable_string = var[Encoding.TYPE]
- if not is_core_window_import_path(importable_string):
- raise WindowNotSupported(importable_string)
- window_cls: type[Window] = import_string(importable_string)
+ if is_core_window_import_path(importable_string):
+ window_cls: type[Window] = import_string(importable_string)
+ else:
+ window_cls = find_registered_custom_window(importable_string)
return window_cls.deserialize(var[Encoding.VAR])
diff --git a/airflow-core/src/airflow/serialization/encoders.py b/airflow-core/src/airflow/serialization/encoders.py
index 59ac29f111e..b31b7c828c6 100644
--- a/airflow-core/src/airflow/serialization/encoders.py
+++ b/airflow-core/src/airflow/serialization/encoders.py
@@ -88,9 +88,9 @@ from airflow.serialization.definitions.deadline import SerializedDeadlineAlert
from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
from airflow.serialization.helpers import (
WaitPolicyNotSupported,
- WindowNotSupported,
find_registered_custom_partition_mapper,
find_registered_custom_timetable,
+ find_registered_custom_window,
is_core_partition_mapper_import_path,
is_core_timetable_import_path,
is_core_wait_policy_import_path,
@@ -695,11 +695,9 @@ def encode_window(var: Window | CoreWindow) -> dict[str, Any]:
"""
Encode a :class:`Window` instance.
- Only built-in ``Window`` subclasses are accepted. Custom subclasses raise
- :class:`WindowNotSupported` so the scheduler never deserializes an
- attacker-controlled import path. If a real need for custom windows arises,
- add a plugin registry mirroring ``partition_mapper`` rather than relaxing
- this check.
+ Custom subclasses must be registered via the ``windows`` plugin attribute;
+ unregistered classes raise :class:`WindowNotSupported` so the scheduler
+ never deserializes an attacker-controlled import path.
The ``BUILTIN_WINDOWS`` fast path maps the SDK classes user code instantiates
(e.g. ``from airflow.sdk import WeekWindow``); after deserialization a
@@ -715,9 +713,12 @@ def encode_window(var: Window | CoreWindow) -> dict[str, Any]:
Encoding.TYPE: importable_string,
Encoding.VAR: _serializer.serialize_window(var),
}
+
qn = qualname(var)
- if not is_core_window_import_path(qn):
- raise WindowNotSupported(qn)
+ if is_core_window_import_path(qn) is False:
+ # This raises if not found.
+ find_registered_custom_window(qn)
+
return {
Encoding.TYPE: qn,
Encoding.VAR: _serializer.serialize_window(var),
diff --git a/airflow-core/src/airflow/serialization/helpers.py b/airflow-core/src/airflow/serialization/helpers.py
index c60fd8baa84..5fc12d23b87 100644
--- a/airflow-core/src/airflow/serialization/helpers.py
+++ b/airflow-core/src/airflow/serialization/helpers.py
@@ -30,6 +30,7 @@ from airflow.configuration import conf
if TYPE_CHECKING:
from airflow.models.deadline import DeadlineReferenceType
from airflow.partition_mappers.base import PartitionMapper
+ from airflow.partition_mappers.window import Window
from airflow.timetables.base import Timetable as CoreTimetable
@@ -200,16 +201,15 @@ def is_core_partition_mapper_import_path(importable_string: str) -> bool:
class WindowNotSupported(ValueError):
- """Raise when serialization encounters a non-built-in ``Window`` subclass."""
+ """Raise when serialization encounters an unregistered ``Window`` subclass."""
def __init__(self, type_string: str) -> None:
self.type_string = type_string
def __str__(self) -> str:
return (
- f"Window class {self.type_string!r} is not a built-in. Custom Window "
- "subclasses are not supported; use one of the built-in "
- "windows under ``airflow.partition_mappers.window``."
+ f"Window class {self.type_string!r} is not registered. Custom Window "
+ "subclasses must be registered via the ``windows`` attribute on an AirflowPlugin."
)
@@ -218,6 +218,16 @@ def is_core_window_import_path(importable_string: str) -> bool:
return importable_string.startswith("airflow.partition_mappers.window.")
+def find_registered_custom_window(importable_string: str) -> type[Window]:
+ """Find a user-defined custom window class registered via a plugin."""
+ from airflow import plugins_manager
+
+ window_classes = plugins_manager.get_windows_plugins()
+ with contextlib.suppress(KeyError):
+ return window_classes[importable_string]
+ raise WindowNotSupported(importable_string)
+
+
class WaitPolicyNotSupported(ValueError):
"""Raise when serialization encounters a non-built-in ``WaitPolicy`` subclass."""
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
index da71a52ddf5..42df1955ba5 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py
@@ -34,11 +34,12 @@ class TestGetPlugins:
# Filters
(
{},
- 15,
+ 16,
[
"InformaticaProviderPlugin",
"MetadataCollectionPlugin",
"OpenLineageProviderPlugin",
+ "business_day_window_plugin",
"databricks_workflow",
"decreasing_priority_weight_strategy_plugin",
"edge_executor",
@@ -55,10 +56,14 @@ class TestGetPlugins:
),
(
{"limit": 3, "offset": 3},
- 15,
- ["databricks_workflow", "decreasing_priority_weight_strategy_plugin", "edge_executor"],
+ 16,
+ [
+ "business_day_window_plugin",
+ "databricks_workflow",
+ "decreasing_priority_weight_strategy_plugin",
+ ],
),
- ({"limit": 1}, 15, ["InformaticaProviderPlugin"]),
+ ({"limit": 1}, 16, ["InformaticaProviderPlugin"]),
],
)
def test_should_respond_200(
@@ -158,10 +163,10 @@ class TestGetPlugins:
plugins_page = body["plugins"]
# Even though limit=7, only 6 valid plugins should come back
- assert len(plugins_page) == 6
+ assert len(plugins_page) == 7
assert "test_plugin_invalid" not in [p["name"] for p in plugins_page]
- assert body["total_entries"] == 15
+ assert body["total_entries"] == 16
@skip_if_force_lowest_dependencies_marker
diff --git a/airflow-core/tests/unit/partition_mappers/test_window.py b/airflow-core/tests/unit/partition_mappers/test_window.py
index b4956761c2c..0b43e06123c 100644
--- a/airflow-core/tests/unit/partition_mappers/test_window.py
+++ b/airflow-core/tests/unit/partition_mappers/test_window.py
@@ -17,9 +17,12 @@
from __future__ import annotations
from datetime import datetime, timedelta
+from unittest import mock
import pytest
+from airflow import plugins_manager
+from airflow._shared.module_loading import qualname
from airflow.partition_mappers.base import RollupMapper
from airflow.partition_mappers.temporal import (
StartOfDayMapper,
@@ -40,7 +43,7 @@ from airflow.partition_mappers.window import (
YearWindow,
)
from airflow.serialization.decoders import decode_partition_mapper, decode_window
-from airflow.serialization.encoders import encode_partition_mapper, encode_window
+from airflow.serialization.encoders import _serializer, encode_partition_mapper, encode_window
from airflow.serialization.enums import Encoding
from airflow.serialization.helpers import WindowNotSupported
@@ -507,22 +510,62 @@ class TestSegmentWindow:
assert frozenset(restored.to_upstream("any")) == frozenset({"us", "eu", "apac"})
+class CustomRegWindow(Window):
+ """A custom Window subclass defined at module scope so it has a stable qualname."""
+
+ def to_upstream(self, decoded_downstream):
+ return [decoded_downstream]
+
+
class TestWindowSerializationGate:
- """``encode_window`` / ``decode_window`` must reject non-built-in Windows.
+ """``encode_window`` / ``decode_window`` accept registered custom Windows and reject unregistered ones.
- Custom Window subclasses are not supported: a tampered serialized Dag
- could otherwise name any importable class and have the scheduler
- ``import_string`` it during deserialization.
+ A tampered serialized Dag naming an unregistered import path is still
+ rejected; only classes that appear on a plugin's ``windows`` attribute
+ are allowed through.
"""
- def test_encode_rejects_custom_window_subclass(self):
- class CustomWindow(Window):
+ @pytest.fixture
+ def window_plugin(self, monkeypatch):
+ """Patch the plugins manager to expose CustomRegWindow under its qualname."""
+ qn = qualname(CustomRegWindow)
+ monkeypatch.setattr(
+ plugins_manager,
+ "get_windows_plugins",
+ lambda: {qn: CustomRegWindow},
+ )
+
+ @pytest.mark.usefixtures("window_plugin")
+ def test_registered_custom_window_round_trips(self):
+ w = CustomRegWindow()
+ restored = decode_window(encode_window(w))
+ assert type(restored) is CustomRegWindow
+ assert restored.to_upstream("x") == ["x"]
+
+ def test_unregistered_custom_window_raises_on_encode(self):
+ class UnregisteredWindow(Window):
def to_upstream(self, decoded_downstream):
return ()
- with pytest.raises(WindowNotSupported, match="CustomWindow"):
- encode_window(CustomWindow())
-
- def test_decode_rejects_non_core_import_path(self):
- with pytest.raises(WindowNotSupported, match="os.system"):
- decode_window({Encoding.TYPE: "os.system", Encoding.VAR: {}})
+ with pytest.raises(WindowNotSupported, match="UnregisteredWindow"):
+ encode_window(UnregisteredWindow())
+
+ def test_unregistered_import_path_raises_on_decode(self):
+ with mock.patch("airflow.serialization.decoders.import_string") as mock_import:
+ with pytest.raises(WindowNotSupported, match="os.system"):
+ decode_window({Encoding.TYPE: "os.system", Encoding.VAR: {}})
+ mock_import.assert_not_called()
+
+ def test_builtin_windows_fast_path_unchanged(self):
+ # The fast-path dict must contain exactly 7 entries, one per builtin SDK class.
+ # Verify that every entry maps to the canonical core import path.
+ expected = {
+ "airflow.partition_mappers.window.HourWindow",
+ "airflow.partition_mappers.window.DayWindow",
+ "airflow.partition_mappers.window.WeekWindow",
+ "airflow.partition_mappers.window.MonthWindow",
+ "airflow.partition_mappers.window.QuarterWindow",
+ "airflow.partition_mappers.window.SegmentWindow",
+ "airflow.partition_mappers.window.YearWindow",
+ }
+ assert set(_serializer.BUILTIN_WINDOWS.values()) == expected
diff --git a/airflow-core/tests/unit/plugins/test_plugins_manager.py b/airflow-core/tests/unit/plugins/test_plugins_manager.py
index 9e0501be6e0..4970b180ad2 100644
--- a/airflow-core/tests/unit/plugins/test_plugins_manager.py
+++ b/airflow-core/tests/unit/plugins/test_plugins_manager.py
@@ -29,6 +29,7 @@ import pytest
from airflow._shared.module_loading import qualname
from airflow.configuration import conf
from airflow.listeners.listener import get_listener_manager
+from airflow.partition_mappers.window import Window
from airflow.plugins_manager import AirflowPlugin
from tests_common.test_utils.config import conf_vars
@@ -79,7 +80,7 @@ class TestPluginsManager:
example_plugins_module="airflow.example_dags.plugins",
)
- assert len(plugins) == 10
+ assert len(plugins) == 11
assert not import_errors
for plugin in plugins:
if "AirflowTestOnLoadPlugin" in str(plugin):
@@ -102,7 +103,7 @@ class TestPluginsManager:
):
plugins, import_errors = plugins_manager._get_plugins()
- assert len(plugins) == 3 # three are loaded from examples
+ assert len(plugins) == 4 # four are loaded from examples
assert len(import_errors) == 1
received_logs = caplog.text
@@ -408,3 +409,27 @@ class TestPluginsManager:
with mock.patch("airflow.plugins_manager._load_plugins_from_plugin_directory", return_value=([], [])):
plugins = plugins_manager._get_plugins()[0]
assert len(plugins) == 6
+
+
+class TestWindowPluginRegistration:
+ """``windows`` plugin attribute surfaces via ``get_windows_plugins()``."""
+
+ def test_windows_attribute_surfaces_via_getter(self):
+ from airflow import plugins_manager
+
+ class MyCustomWindow(Window):
+ name = "test_window_plugin"
+
+ def to_upstream(self, decoded_downstream):
+ return [decoded_downstream]
+
+ class MyWindowPlugin(AirflowPlugin):
+ name = "test_window_plugin"
+ windows = [MyCustomWindow]
+
+ with mock_plugin_manager(plugins=[MyWindowPlugin()]):
+ plugins_manager.get_windows_plugins.cache_clear()
+ registered = plugins_manager.get_windows_plugins()
+
+ assert qualname(MyCustomWindow) in registered
+ assert registered[qualname(MyCustomWindow)] is MyCustomWindow
diff --git a/shared/plugins_manager/src/airflow_shared/plugins_manager/plugins_manager.py b/shared/plugins_manager/src/airflow_shared/plugins_manager/plugins_manager.py
index 17739428013..e8dea0d1294 100644
--- a/shared/plugins_manager/src/airflow_shared/plugins_manager/plugins_manager.py
+++ b/shared/plugins_manager/src/airflow_shared/plugins_manager/plugins_manager.py
@@ -122,6 +122,9 @@ class AirflowPlugin:
# A list of partition mapper classes that can be used for Dag scheduling.
partition_mappers: list[Any] = []
+ # A list of window classes that can be used for Dag scheduling.
+ windows: list[Any] = []
+
# A list of deadline reference classes that can be used as custom deadlines in Dags.
deadline_references: list[Any] = []