This is an automated email from the ASF dual-hosted git repository.
o-nikolas 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 836a5565c58 Scope plugin scheduling classes to their team in
multi-team mode (#73569)
836a5565c58 is described below
commit 836a5565c58afd0a9267827b6692c955bf1ccb7e
Author: Niko Oliveira <[email protected]>
AuthorDate: Tue Sep 29 11:22:03 2026 -0700
Scope plugin scheduling classes to their team in multi-team mode (#73569)
A team-scoped plugin's timetables, partition mappers, windows, deadline
references and priority weight strategies were usable by any Dag. Dags are
now checked when their parsing results are stored: a Dag that uses another
team's class gets an import error instead of being stored.
---
.../src/airflow/dag_processing/collection.py | 104 ++++++
airflow-core/src/airflow/plugins_manager.py | 42 +++
.../tests/unit/dag_processing/test_collection.py | 347 ++++++++++++++++++++-
.../tests/unit/plugins/test_plugins_manager.py | 61 ++++
.../src/tests_common/test_utils/mock_plugins.py | 42 ++-
5 files changed, 578 insertions(+), 18 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/collection.py
b/airflow-core/src/airflow/dag_processing/collection.py
index 70c95f8eedf..33b1b443931 100644
--- a/airflow-core/src/airflow/dag_processing/collection.py
+++ b/airflow-core/src/airflow/dag_processing/collection.py
@@ -50,11 +50,13 @@ from airflow.models.asset import (
TaskOutletAssetReference,
)
from airflow.models.dag import DagModel, DagOwnerAttributes, DagTag
+from airflow.models.dagbundle import DagBundleModel
from airflow.models.dagrun import DagRun
from airflow.models.dagwarning import DagWarning, DagWarningType
from airflow.models.errors import ParseImportError
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.trigger import Trigger
+from airflow.plugins_manager import get_scheduling_class_teams
from airflow.serialization.definitions.assets import (
SerializedAsset,
SerializedAssetAlias,
@@ -475,6 +477,101 @@ def _update_import_errors(
)
+def _iter_serialized_class_names(data: Any) -> Iterator[str]:
+ """
+ Yield every string in a serialized Dag, which includes every class name it
records.
+
+ Timetables, partition mappers and windows are encoded as ``{"__type":
<qualname>, ...}``,
+ a custom deadline reference carries ``"__class_path"``, and a task's
``weight_rule`` is
+ stored as the strategy's qualname. Yielding every string rather than
reading known keys
+ means it does not matter where the encoders put a name; ordinary strings
simply match
+ nothing in the plugin map.
+
+ This is what ``DagSerialization.to_dict`` returns, before it is dumped to
JSON, so it
+ holds tuples as well as lists (a partition mapper config is a list of
pairs), and names
+ can sit in keys as well as values. None of it is cyclic.
+ """
+ stack = [data]
+ while stack:
+ item = stack.pop()
+ if isinstance(item, str):
+ yield item
+ elif isinstance(item, dict):
+ stack.extend(item.keys())
+ stack.extend(item.values())
+ elif isinstance(item, (list, tuple, set, frozenset)):
+ stack.extend(item)
+
+
+def _reject_other_teams_plugin_classes(
+ bundle_name: str,
+ dags: Collection[LazyDeserializedDAG],
+ import_errors: dict[tuple[str, str], str],
+ *,
+ session: Session,
+) -> list[LazyDeserializedDAG]:
+ """
+ Drop Dags that name a scheduling class belonging to another team's plugin.
+
+ Timetables, partition mappers, windows, deadline references and priority
weight
+ strategies are named by the Dag itself, with no team-aware lookup in
between, so this is
+ the only thing keeping a team-scoped plugin's scheduling classes to that
team.
+
+ It runs here, on the data about to be stored, because that is exactly what
the scheduler
+ will resolve, and because on the Dag processor path this runs in the
manager, which never
+ executes Dag code. Configuration read in the process that imported the Dag
file cannot be
+ trusted: the file could have changed it.
+
+ A partition mapper a timetable picks inside ``get_partition_mapper()`` is
not covered,
+ because nothing decides it until the timetable runs.
+ """
+ if not conf.getboolean("core", "multi_team"):
+ return list(dags)
+
+ # A class registered by any global plugin is available to every Dag, so
only classes
+ # that every registering plugin scoped to a team can be refused.
+ restricted: dict[str, frozenset[str]] = {
+ name: frozenset(team for team in teams if team is not None)
+ for name, teams in get_scheduling_class_teams().items()
+ if None not in teams
+ }
+ if not restricted:
+ return list(dags)
+
+ accepted: list[LazyDeserializedDAG] = []
+ dag_team_name: str | None = None
+ team_looked_up = False
+ for dag in dags:
+ used = {
+ name: restricted[name] for name in
_iter_serialized_class_names(dag.data) if name in restricted
+ }
+ if used and not team_looked_up:
+ dag_team_name = DagBundleModel.get_team_name(bundle_name,
session=session)
+ team_looked_up = True
+ refused = {name: teams for name, teams in used.items() if
dag_team_name not in teams}
+ if not refused:
+ accepted.append(dag)
+ continue
+
+ name, owning_teams = min(refused.items())
+ owners = ", ".join(sorted(owning_teams))
+ belongs_to = f"team '{dag_team_name}'" if dag_team_name else "no team"
+ log.warning(
+ "Refusing Dag that uses another team's plugin class",
+ dag_id=dag.dag_id,
+ class_name=name,
+ owning_teams=sorted(owning_teams),
+ dag_team=dag_team_name,
+ )
+ import_errors[(bundle_name, dag.relative_fileloc)] = (
+ f"Dag '{dag.dag_id}' uses {name}, which is provided by a plugin
belonging to {owners}. "
+ f"This Dag belongs to {belongs_to}, so it cannot use it. Move the
Dag into a bundle "
+ f"owned by {owners}, or have the plugin provide the class globally
instead of for a "
+ "single team."
+ )
+ return accepted
+
+
def update_dag_parsing_results_in_db(
bundle_name: str,
bundle_version: str | None,
@@ -513,6 +610,13 @@ def update_dag_parsing_results_in_db(
If None, will be inferred from dags and import_errors. Passing this
explicitly ensures that
import errors are cleared for files that were parsed but no longer
contain DAGs.
"""
+ accepted = _reject_other_teams_plugin_classes(bundle_name, dags,
import_errors, session=session)
+ if len(accepted) != len(dags):
+ # A rejected Dag may have no ``dag`` row yet, and dag_warning has a
foreign key to it.
+ rejected_ids = {dag.dag_id for dag in dags} - {dag.dag_id for dag in
accepted}
+ warnings = {warning for warning in warnings if warning.dag_id not in
rejected_ids}
+ dags = accepted
+
# Retry 'DAG.bulk_write_to_db' & 'SerializedDagModel.bulk_sync_to_db' in
case
# of any Operational Errors
# In case of failures, provide_session handles rollback
diff --git a/airflow-core/src/airflow/plugins_manager.py
b/airflow-core/src/airflow/plugins_manager.py
index edf8db34d5b..8b2ea644e0b 100644
--- a/airflow-core/src/airflow/plugins_manager.py
+++ b/airflow-core/src/airflow/plugins_manager.py
@@ -44,6 +44,10 @@ from airflow._shared.plugins_manager import (
is_valid_plugin,
)
from airflow.configuration import conf
+from airflow.serialization.helpers import (
+ is_core_partition_mapper_import_path,
+ is_core_timetable_import_path,
+)
if TYPE_CHECKING:
from airflow.listeners.listener import ListenerManager
@@ -519,6 +523,44 @@ def is_extra_link_visible_to_team(link: Any, team_name:
str | None) -> bool:
return team_name in link_teams
+@cache
+def get_scheduling_class_teams() -> dict[str, frozenset[str | None]]:
+ """
+ Map the qualname of every plugin-registered scheduling class to the teams
that registered it.
+
+ Covers timetables, partition mappers, windows, deadline references and
priority weight
+ strategies: the registries a Dag names directly, with no team-aware lookup
in between.
+
+ Keyed by qualname because that is what a serialized Dag records and what
the scheduler
+ resolves through ``get_timetables_plugins()`` and its siblings. Class
identity is not
+ stable enough to key on: the plugin loader executes a plugin file again
under its own
+ module entry, so a Dag importing a class from that file can hold a
different class
+ object with the same qualname as the one that was registered.
+
+ A qualname registered by several plugins maps to all of their teams, and
is then
+ resolved least restrictively.
+
+ Airflow's own timetables, partition mappers and windows are left out even
if a plugin
+ lists them: the decoder imports anything under those core paths directly
and never
+ consults plugins, so a plugin cannot own them.
+ """
+ teams: dict[str, set[str | None]] = {}
+ for plugin in _get_plugins()[0]:
+ for scheduling_class in (
+ *plugin.timetables,
+ *plugin.partition_mappers,
+ *plugin.windows,
+ *plugin.deadline_references,
+ *plugin.priority_weight_strategies,
+ ):
+ name = qualname(scheduling_class)
+ # The partition mapper prefix also covers core windows.
+ if is_core_timetable_import_path(name) or
is_core_partition_mapper_import_path(name):
+ continue
+ teams.setdefault(name, set()).add(plugin.team_name)
+ return {name: frozenset(team_names) for name, team_names in teams.items()}
+
+
@cache
def get_timetables_plugins() -> dict[str, type[Timetable]]:
"""Collect and get timetable classes registered by plugins."""
diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py
b/airflow-core/tests/unit/dag_processing/test_collection.py
index be36290abce..cc331423e3f 100644
--- a/airflow-core/tests/unit/dag_processing/test_collection.py
+++ b/airflow-core/tests/unit/dag_processing/test_collection.py
@@ -18,7 +18,11 @@
from __future__ import annotations
+import importlib
import logging
+import os
+import sys
+import textwrap
import warnings
from collections.abc import Generator
from datetime import timedelta
@@ -31,6 +35,8 @@ from sqlalchemy import delete, event, func, inspect as
sa_inspect, select
from sqlalchemy.exc import OperationalError, SAWarning
import airflow.dag_processing.collection
+from airflow import plugins_manager
+from airflow._shared.module_loading import qualname
from airflow._shared.timezones import timezone as tz
from airflow.configuration import conf
from airflow.dag_processing.collection import (
@@ -42,6 +48,9 @@ from airflow.dag_processing.collection import (
_update_import_errors,
update_dag_parsing_results_in_db,
)
+from airflow.example_dags.plugins.business_day_window import BusinessDayWindow
+from airflow.example_dags.plugins.custom_partition_mapper import
PrefixStripMapper
+from airflow.example_dags.plugins.workday import AfterWorkdayTimetable
from airflow.exceptions import SerializationError
from airflow.models import DagModel, DagRun
from airflow.models.asset import (
@@ -57,27 +66,36 @@ from airflow.models.errors import ParseImportError
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.trigger import Trigger
from airflow.partition_mappers.base import RollupMapper
-from airflow.partition_mappers.temporal import StartOfDayMapper
+from airflow.partition_mappers.chain import ChainMapper
+from airflow.partition_mappers.identity import IdentityMapper
+from airflow.partition_mappers.temporal import StartOfDayMapper,
StartOfMonthMapper
from airflow.partition_mappers.window import DayWindow
+from airflow.plugins_manager import AirflowPlugin
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.triggers.file import FileDeleteTrigger
from airflow.sdk import DAG, Asset, AssetAlias, AssetAll, AssetWatcher
-from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
+from airflow.sdk.definitions.deadline import AsyncCallback,
BaseDeadlineReference, DeadlineAlert
+from airflow.sdk.definitions.timetables.assets import AssetOrTimeSchedule,
PartitionedAssetTimetable
from airflow.serialization.definitions.assets import SerializedAsset
from airflow.serialization.encoders import encode_trigger,
ensure_serialized_asset
from airflow.serialization.serialized_objects import LazyDeserializedDAG
from airflow.timetables.simple import PartitionedAtRuntime
+from airflow.timetables.trigger import CronTriggerTimetable
from airflow.triggers.base import BaseEventTrigger
from airflow.utils.types import DagRunType
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
clear_db_assets,
+ clear_db_dag_bundles,
clear_db_dags,
clear_db_import_errors,
clear_db_serialized_dags,
+ clear_db_teams,
clear_db_triggers,
)
+from tests_common.test_utils.mock_plugins import mock_plugin_manager
+from unit.plugins.priority_weight_strategy import
StaticTestPriorityWeightStrategy
if TYPE_CHECKING:
from kgb import SpyAgency
@@ -1680,3 +1698,328 @@ class TestPartitionMapperInfoSync:
dag_model = session.get(DagModel, "non_partitioned_dag")
assert dag_model.partition_mapper_info == []
assert dag_model.has_rollup_mappers is False
+
+
+class TeamDeadlineReference(BaseDeadlineReference):
+ """A deadline reference a team-scoped plugin ships; Airflow has no example
one to reuse."""
+
+ def _evaluate_with(self, *, session, **kwargs):
+ raise NotImplementedError
+
+
+async def _deadline_callback():
+ raise NotImplementedError
+
+
+def _nested_chain_mapper(depth):
+ mapper = PrefixStripMapper("eu")
+ for _ in range(depth):
+ mapper = ChainMapper(mapper, IdentityMapper())
+ return mapper
+
+
+def _dag_kwargs_using(case):
+ """Return the Dag keyword arguments that make it use the plugin class, as
a Dag author would."""
+ if case == "timetable":
+ return {"schedule": AfterWorkdayTimetable()}
+ if case == "timetable-in-asset-or-time":
+ return {"schedule":
AssetOrTimeSchedule(timetable=AfterWorkdayTimetable(), assets=[Asset("a")])}
+ if case == "default-partition-mapper":
+ return {
+ "schedule": PartitionedAssetTimetable(
+ assets=Asset("a"),
default_partition_mapper=PrefixStripMapper("eu")
+ )
+ }
+ if case == "partition-mapper-deep-in-chain":
+ return {
+ "schedule": PartitionedAssetTimetable(
+ assets=Asset("a"), partition_mapper_config={Asset("a"):
_nested_chain_mapper(6)}
+ )
+ }
+ if case == "window-in-rollup-mapper":
+ return {
+ "schedule": PartitionedAssetTimetable(
+ assets=Asset("a"),
+ partition_mapper_config={
+ Asset("a"): RollupMapper(window=BusinessDayWindow(),
upstream_mapper=StartOfMonthMapper())
+ },
+ )
+ }
+ if case == "deadline-reference":
+ return {
+ "deadline": DeadlineAlert(
+ reference=TeamDeadlineReference(),
+ interval=timedelta(hours=1),
+ callback=AsyncCallback(_deadline_callback),
+ )
+ }
+ raise ValueError(case)
+
+
+# Each case: the registry the plugin fills, the class it registers, and how
the Dag uses it.
+SCHEDULING_CLASS_USES = [
+ pytest.param("timetables", AfterWorkdayTimetable, "timetable",
id="timetable"),
+ pytest.param(
+ "timetables", AfterWorkdayTimetable, "timetable-in-asset-or-time",
id="timetable-in-asset-or-time"
+ ),
+ pytest.param(
+ "partition_mappers", PrefixStripMapper, "default-partition-mapper",
id="default-partition-mapper"
+ ),
+ pytest.param(
+ "partition_mappers",
+ PrefixStripMapper,
+ "partition-mapper-deep-in-chain",
+ id="partition-mapper-deep-in-chain",
+ ),
+ pytest.param("windows", BusinessDayWindow, "window-in-rollup-mapper",
id="window-in-rollup-mapper"),
+ pytest.param("deadline_references", TeamDeadlineReference,
"deadline-reference", id="deadline-reference"),
+]
+
+
[email protected]_test
+class TestRejectOtherTeamsPluginClasses:
+ """A team-scoped plugin's scheduling classes may only be stored for that
team's Dags."""
+
+ @pytest.fixture(autouse=True)
+ def _clean(self):
+ yield
+ clear_db_serialized_dags()
+ clear_db_dags()
+ clear_db_import_errors()
+ clear_db_dag_bundles()
+ clear_db_teams()
+
+ @pytest.fixture
+ def bundle(self, testing_team, session):
+ """Return a factory creating the "team_bundle" bundle, owned by the
given team or none."""
+
+ def create(owned_by_team: bool) -> str:
+ bundle = DagBundleModel(name="team_bundle")
+ if owned_by_team:
+ bundle.teams.append(testing_team)
+ session.add(bundle)
+ session.flush()
+ return bundle.name
+
+ return create
+
+ @staticmethod
+ def _plugin(team_name, registry, scheduling_class,
name="scheduling_plugin"):
+ plugin = AirflowPlugin()
+ plugin.name = name
+ plugin.team_name = team_name
+ setattr(plugin, registry, [scheduling_class])
+ return plugin
+
+ @staticmethod
+ def _serialized(dag_id="team_dag", **dag_kwargs):
+ with DAG(dag_id, **{"schedule": None, **dag_kwargs}) as dag:
+ EmptyOperator(task_id="t")
+ dag.relative_fileloc = f"{dag_id}.py"
+ return LazyDeserializedDAG.from_dag(dag)
+
+ @staticmethod
+ def _store(bundle_name, dags, session, warnings=frozenset()):
+ import_errors: dict[tuple[str, str], str] = {}
+ update_dag_parsing_results_in_db(
+ bundle_name=bundle_name,
+ bundle_version=None,
+ dags=dags,
+ import_errors=import_errors,
+ parse_duration=None,
+ warnings=set(warnings),
+ session=session,
+ )
+ stored = set(session.scalars(select(SerializedDagModel.dag_id)))
+ errors = {
+ (e.bundle_name, e.filename): e.stacktrace for e in
session.scalars(select(ParseImportError))
+ }
+ return stored, errors
+
+ @conf_vars({("core", "multi_team"): "True"})
+ @pytest.mark.parametrize(("registry", "scheduling_class", "usage"),
SCHEDULING_CLASS_USES)
+ @pytest.mark.parametrize(
+ ("plugin_team", "dag_owned_by_team", "allowed"),
+ [
+ pytest.param("testing", True, True, id="owning-team"),
+ pytest.param("other_team", True, False, id="other-team"),
+ pytest.param("testing", False, False, id="teamless"),
+ ],
+ )
+ def test_team_class_is_only_stored_for_its_team(
+ self, bundle, session, plugin_team, dag_owned_by_team, allowed,
registry, scheduling_class, usage
+ ):
+ bundle_name = bundle(dag_owned_by_team)
+ with mock_plugin_manager(plugins=[self._plugin(plugin_team, registry,
scheduling_class)]):
+ dag = self._serialized(**_dag_kwargs_using(usage))
+ stored, errors = self._store(bundle_name, [dag], session)
+
+ if allowed:
+ assert stored == {"team_dag"}
+ assert errors == {}
+ else:
+ assert stored == set()
+ assert list(errors) == [(bundle_name, "team_dag.py")]
+ assert f"belonging to {plugin_team}" in errors[(bundle_name,
"team_dag.py")]
+
+ @conf_vars({("core", "multi_team"): "True"})
+ @pytest.mark.parametrize(
+ "weight_rule",
+ [StaticTestPriorityWeightStrategy(),
qualname(StaticTestPriorityWeightStrategy)],
+ ids=["instance", "dotted-path"],
+ )
+ def test_weight_rule_is_checked_in_both_spellings(self, bundle, session,
weight_rule):
+ bundle_name = bundle(True)
+ plugin = self._plugin("other_team", "priority_weight_strategies",
StaticTestPriorityWeightStrategy)
+ with mock_plugin_manager(plugins=[plugin]):
+ with DAG("team_dag", schedule=None) as dag:
+ EmptyOperator(task_id="t", weight_rule=weight_rule)
+ dag.relative_fileloc = "team_dag.py"
+ stored, errors = self._store(bundle_name,
[LazyDeserializedDAG.from_dag(dag)], session)
+
+ assert stored == set()
+ assert "belonging to other_team" in errors[(bundle_name,
"team_dag.py")]
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_only_the_offending_dag_is_dropped(self, bundle, session):
+ bundle_name = bundle(True)
+ with mock_plugin_manager(plugins=[self._plugin("other_team",
"timetables", AfterWorkdayTimetable)]):
+ dags = [
+ self._serialized("rejected", schedule=AfterWorkdayTimetable()),
+ self._serialized("accepted"),
+ ]
+ stored, errors = self._store(bundle_name, dags, session)
+
+ assert stored == {"accepted"}
+ assert list(errors) == [(bundle_name, "rejected.py")]
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_warning_for_a_rejected_new_dag_is_dropped(self, bundle, session):
+ """
+ The stability check warns about every Dag in a file, including one
being rejected.
+
+ A new Dag has no ``dag`` row to hang that warning on, so storing it
would break the
+ foreign key and fail the whole write.
+ """
+ bundle_name = bundle(True)
+ warnings = {
+ DagWarning("rejected", DagWarningType.RUNTIME_VARYING_VALUE.value,
"datetime.now() in args"),
+ DagWarning("accepted", DagWarningType.RUNTIME_VARYING_VALUE.value,
"datetime.now() in args"),
+ }
+ with mock_plugin_manager(plugins=[self._plugin("other_team",
"timetables", AfterWorkdayTimetable)]):
+ dags = [
+ self._serialized("rejected", schedule=AfterWorkdayTimetable()),
+ self._serialized("accepted"),
+ ]
+ stored, errors = self._store(bundle_name, dags, session,
warnings=warnings)
+
+ assert stored == {"accepted"}
+ assert list(errors) == [(bundle_name, "rejected.py")]
+ assert set(session.scalars(select(DagWarning.dag_id))) == {"accepted"}
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_error_is_a_plain_message(self, bundle, session):
+ """The UI shows this as-is, so it must read as an explanation, not a
traceback."""
+ bundle_name = bundle(True)
+ with mock_plugin_manager(plugins=[self._plugin("other_team",
"timetables", AfterWorkdayTimetable)]):
+ stored, errors = self._store(
+ bundle_name,
[self._serialized(schedule=AfterWorkdayTimetable())], session
+ )
+
+ assert errors[(bundle_name, "team_dag.py")] == (
+ f"Dag 'team_dag' uses {qualname(AfterWorkdayTimetable)}, which is
provided by a plugin "
+ "belonging to other_team. This Dag belongs to team 'testing', so
it cannot use it. "
+ "Move the Dag into a bundle owned by other_team, or have the
plugin provide the class "
+ "globally instead of for a single team."
+ )
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_class_also_registered_globally_is_available_to_every_dag(self,
bundle, session):
+ bundle_name = bundle(True)
+ plugins = [
+ self._plugin(team, "timetables", AfterWorkdayTimetable,
name=f"plugin_{i}")
+ for i, team in enumerate(["other_team", None])
+ ]
+ with mock_plugin_manager(plugins=plugins):
+ stored, errors = self._store(
+ bundle_name,
[self._serialized(schedule=AfterWorkdayTimetable())], session
+ )
+
+ assert stored == {"team_dag"}
+ assert errors == {}
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_airflow_class_listed_by_a_team_plugin_stays_available(self,
bundle, session):
+ """
+ Every team's cron Dags keep working even if one team's plugin lists
the cron timetable.
+
+ The timetable is explicit: a cron string can serialize to a different
class, depending on
+ ``create_cron_data_intervals``, which would leave the plugin's class
out of the Dag.
+ """
+ bundle_name = bundle(True)
+ with mock_plugin_manager(plugins=[self._plugin("other_team",
"timetables", CronTriggerTimetable)]):
+ stored, errors = self._store(
+ bundle_name,
+ [self._serialized(schedule=CronTriggerTimetable("0 0 * * *",
timezone="UTC"))],
+ session,
+ )
+
+ assert stored == {"team_dag"}
+ assert errors == {}
+
+ def test_nothing_is_rejected_when_multi_team_is_off(self, bundle, session):
+ bundle_name = bundle(True)
+ with mock_plugin_manager(plugins=[self._plugin("other_team",
"timetables", AfterWorkdayTimetable)]):
+ stored, errors = self._store(
+ bundle_name,
[self._serialized(schedule=AfterWorkdayTimetable())], session
+ )
+
+ assert stored == {"team_dag"}
+ assert errors == {}
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_class_reloaded_by_the_plugin_loader_is_still_recognised(
+ self, bundle, session, tmp_path, monkeypatch, request
+ ):
+ """
+ A Dag importing from a plugin file holds a different class from the
one registered.
+
+ The plugin loader executes the file again under its own module entry,
so the two
+ classes share a qualname but not an identity.
+ """
+ (tmp_path / "workday.py").write_text(
+ textwrap.dedent(
+ """\
+ from airflow.plugins_manager import AirflowPlugin
+ from airflow.timetables.simple import NullTimetable
+
+
+ class WorkdayTimetable(NullTimetable):
+ pass
+
+
+ class WorkdayPlugin(AirflowPlugin):
+ name = "workday"
+ team_name = "other_team"
+ timetables = [WorkdayTimetable]
+ """
+ )
+ )
+ monkeypatch.syspath_prepend(os.fspath(tmp_path))
+ # Both the import below and the plugin loader put a "workday" module
in sys.modules, and
+ # monkeypatch would restore the loader's at teardown, so remove it
outright instead.
+ sys.modules.pop("workday", None)
+ request.addfinalizer(lambda: sys.modules.pop("workday", None))
+ dag_side_class = importlib.import_module("workday").WorkdayTimetable
+ plugins, import_errors =
plugins_manager._load_plugins_from_plugin_directory(
+ plugins_folder=os.fspath(tmp_path)
+ )
+ assert not import_errors
+ assert plugins[0].timetables[0] is not dag_side_class
+
+ bundle_name = bundle(True)
+ with mock_plugin_manager(plugins=plugins):
+ stored, errors = self._store(bundle_name,
[self._serialized(schedule=dag_side_class())], session)
+
+ assert stored == set()
+ assert "belonging to other_team" in errors[(bundle_name,
"team_dag.py")]
diff --git a/airflow-core/tests/unit/plugins/test_plugins_manager.py
b/airflow-core/tests/unit/plugins/test_plugins_manager.py
index bf6d4f00f4c..b1a883e9282 100644
--- a/airflow-core/tests/unit/plugins/test_plugins_manager.py
+++ b/airflow-core/tests/unit/plugins/test_plugins_manager.py
@@ -645,6 +645,67 @@ class TestPluginTeamName:
assert info_by_name["team_plugin"]["team_name"] == "team_a"
+class TestGetSchedulingClassTeams:
+ @staticmethod
+ def _plugin(team_name, **registries):
+ plugin = AirflowPlugin()
+ plugin.name = f"plugin_{team_name}"
+ plugin.team_name = team_name
+ for registry, classes in registries.items():
+ setattr(plugin, registry, classes)
+ return plugin
+
+ def test_maps_each_registry_by_qualname(self):
+ from airflow.example_dags.plugins.business_day_window import
BusinessDayWindow
+ from airflow.example_dags.plugins.custom_partition_mapper import
PrefixStripMapper
+ from airflow.example_dags.plugins.decreasing_priority_weight_strategy
import (
+ DecreasingPriorityStrategy,
+ )
+ from airflow.example_dags.plugins.workday import AfterWorkdayTimetable
+
+ plugin = self._plugin(
+ "team_a",
+ timetables=[AfterWorkdayTimetable],
+ partition_mappers=[PrefixStripMapper],
+ windows=[BusinessDayWindow],
+ priority_weight_strategies=[DecreasingPriorityStrategy],
+ )
+ with mock_plugin_manager(plugins=[plugin]):
+ assert plugins_manager.get_scheduling_class_teams() == {
+ qualname(cls): frozenset({"team_a"})
+ for cls in (
+ AfterWorkdayTimetable,
+ PrefixStripMapper,
+ BusinessDayWindow,
+ DecreasingPriorityStrategy,
+ )
+ }
+
+ def test_class_registered_by_several_plugins_maps_to_all_their_teams(self):
+ from airflow.example_dags.plugins.workday import AfterWorkdayTimetable
+
+ plugins = [self._plugin(team, timetables=[AfterWorkdayTimetable]) for
team in ("team_a", None)]
+ with mock_plugin_manager(plugins=plugins):
+ assert plugins_manager.get_scheduling_class_teams() == {
+ qualname(AfterWorkdayTimetable): frozenset({"team_a", None})
+ }
+
+ def test_airflow_classes_are_left_out(self):
+ """The decoder imports these directly, so no plugin can own them."""
+ from airflow.partition_mappers.temporal import StartOfDayMapper
+ from airflow.partition_mappers.window import DayWindow
+ from airflow.timetables.trigger import CronTriggerTimetable
+
+ plugin = self._plugin(
+ "team_a",
+ timetables=[CronTriggerTimetable],
+ partition_mappers=[StartOfDayMapper],
+ windows=[DayWindow],
+ )
+ with mock_plugin_manager(plugins=[plugin]):
+ assert plugins_manager.get_scheduling_class_teams() == {}
+
+
class TestValidatePluginTeams:
"""``validate_plugin_teams`` startup validation."""
diff --git a/devel-common/src/tests_common/test_utils/mock_plugins.py
b/devel-common/src/tests_common/test_utils/mock_plugins.py
index cc100747dd0..b33831bfbce 100644
--- a/devel-common/src/tests_common/test_utils/mock_plugins.py
+++ b/devel-common/src/tests_common/test_utils/mock_plugins.py
@@ -19,7 +19,7 @@ from __future__ import annotations
from contextlib import ExitStack, contextmanager
from unittest import mock
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS, AIRFLOW_V_3_3_PLUS
PLUGINS_MANAGER_NULLABLE_ATTRIBUTES_V3_0 = [
"plugins",
@@ -83,24 +83,34 @@ def mock_plugin_manager(plugins=None, **kwargs):
# Handle plugins specially
with ExitStack() as exit_stack:
if AIRFLOW_V_3_2_PLUS:
- # Always start the block with an non-initialized plugins, so
ensure_plugins_loaded runs.
from airflow import plugins_manager
from airflow.sdk import plugins_manager as sdk_plugins_manager
- plugins_manager._get_plugins.cache_clear()
- plugins_manager._get_ui_plugins.cache_clear()
- plugins_manager.get_ui_translations.cache_clear()
- plugins_manager.get_flask_plugins.cache_clear()
- plugins_manager.get_fastapi_plugins.cache_clear()
- plugins_manager._get_extra_operators_links_plugins.cache_clear()
- plugins_manager._get_extra_link_class_teams.cache_clear()
- plugins_manager.get_timetables_plugins.cache_clear()
- plugins_manager.integrate_macros_plugins.cache_clear()
- plugins_manager.get_priority_weight_strategy_plugins.cache_clear()
-
- sdk_plugins_manager.integrate_macros_plugins.cache_clear()
- sdk_plugins_manager.get_macro_plugin_teams.cache_clear()
- sdk_plugins_manager.get_hook_lineage_readers_plugins.cache_clear()
+ def clear_plugin_caches():
+ plugins_manager._get_plugins.cache_clear()
+ plugins_manager._get_ui_plugins.cache_clear()
+ plugins_manager.get_ui_translations.cache_clear()
+ plugins_manager.get_flask_plugins.cache_clear()
+ plugins_manager.get_fastapi_plugins.cache_clear()
+
plugins_manager._get_extra_operators_links_plugins.cache_clear()
+ plugins_manager._get_extra_link_class_teams.cache_clear()
+ plugins_manager.get_scheduling_class_teams.cache_clear()
+ plugins_manager.get_timetables_plugins.cache_clear()
+ plugins_manager.get_partition_mapper_plugins.cache_clear()
+ if AIRFLOW_V_3_3_PLUS:
+ plugins_manager.get_windows_plugins.cache_clear()
+
plugins_manager.get_deadline_references_plugins.cache_clear()
+ plugins_manager.integrate_macros_plugins.cache_clear()
+
plugins_manager.get_priority_weight_strategy_plugins.cache_clear()
+
+ sdk_plugins_manager.integrate_macros_plugins.cache_clear()
+ sdk_plugins_manager.get_macro_plugin_teams.cache_clear()
+
sdk_plugins_manager.get_hook_lineage_readers_plugins.cache_clear()
+
+ # Start the block with no plugins loaded, so ensure_plugins_loaded
runs, and clear again
+ # on the way out so nothing computed from the mocked plugins
outlives the block.
+ clear_plugin_caches()
+ exit_stack.callback(clear_plugin_caches)
if plugins is not None or "import_errors" in kwargs:
exit_stack.enter_context(
