Scope plugin operator extra links to their team in multi-team mode (#72231)

A team-scoped plugin's operator extra links were rendered on every task
instance, so one team's links appeared on other teams' Dags and pointed
users at systems they may have no access to.

Ownership is tracked per link class rather than per link name: link
instances are unhashable and compare equal to each other, and a plugin
link may deliberately share a name with an operator's own link, so a
name key would make the operator's link look team-owned. A class
registered by both a global and a team-scoped plugin stays global.
This commit is contained in:
Niko Oliveira authored and GitHub committed 2026-09-14 19:13:49 -07:00
1 parent 4c3ec9cbae
commit b1df686509
5 files changed
+314 -2

No files matched your search

@@ -17,23 +17,42 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from fastapi import Depends, HTTPException, status
from sqlalchemy.sql import select
from airflow import plugins_manager
from airflow.api_fastapi.common.dagbag import DagBagDep, get_dag_for_run_or_latest_version
from airflow.api_fastapi.common.db.common import SessionDep
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.api_fastapi.core_api.datamodels.extra_links import ExtraLinkCollectionResponse
from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
from airflow.api_fastapi.core_api.security import DagAccessEntity, requires_access_dag
from airflow.configuration import conf
from airflow.exceptions import TaskNotFound
from airflow.models import DagRun
from airflow.models.dag import DagModel
if TYPE_CHECKING:
from airflow.serialization.serialized_objects import SerializedOperator
extra_links_router = AirflowRouter(
tags=["Extra Links"], prefix="/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances/{task_id}/links"
)
def _find_operator_link(task: SerializedOperator, link_name: str) -> Any:
"""
Resolve a link name to the link object that will render it.
Mirrors the lookup order of ``get_extra_links`` so the object inspected for team
ownership is the one actually used, which matters when an operator link and a
plugin link share a name.
"""
return task.operator_extra_link_dict.get(link_name) or task.global_operator_extra_link_dict.get(link_name)
@extra_links_router.get(
"",
responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
@@ -99,9 +118,20 @@ def get_extra_links(
else:
ti_for_links = ti
link_names: list[str] = task.extra_links
if conf.getboolean("core", "multi_team"):
dag_team_name = DagModel.get_team_name(dag_id)
link_names = [
link_name
for link_name in link_names
if plugins_manager.is_extra_link_visible_to_team(
_find_operator_link(task, link_name), dag_team_name
)
]
all_extra_link_pairs = (
(link_name, task.get_extra_links(ti_for_links, link_name))
for link_name in task.extra_links # type: ignore[arg-type]
for link_name in link_names # type: ignore[arg-type]
)
all_extra_links = {link_name: link_url or None for link_name, link_url in sorted(all_extra_link_pairs)}
@@ -479,6 +479,46 @@ def get_operator_extra_links() -> list[Any]:
return _get_extra_operators_links_plugins()[1]
@cache
def _get_extra_link_class_teams() -> dict[type, frozenset[str | None]]:
"""
Map every plugin-registered extra link class to the teams that registered it.
Keyed by class because neither of the alternatives works: ``BaseOperatorLink`` sets
``__hash__ = None`` and compares equal across instances, and two distinct link
classes may share a ``name`` (plugin links deliberately override operator links of
the same name), so a name key would conflate them.
A class registered by several plugins maps to all of their teams, which
:func:`is_extra_link_visible_to_team` then resolves least restrictively.
"""
teams: dict[type, set[str | None]] = {}
for plugin in _get_plugins()[0]:
for link in (*plugin.global_operator_extra_links, *plugin.operator_extra_links):
teams.setdefault(type(link), set()).add(plugin.team_name)
return {link_class: frozenset(team_names) for link_class, team_names in teams.items()}
def is_extra_link_visible_to_team(link: Any, team_name: str | None) -> bool:
"""
Whether ``link`` should be shown on a task instance belonging to ``team_name``.
A team-scoped plugin's links are shown only on that team's task instances, so they
appear neither on another team's Dags nor on teamless (global) ones. Links from
global plugins, and links the operator defines itself, stay visible everywhere.
:param link: The operator link object, whose class identifies the registering plugin.
:param team_name: Team owning the Dag the link would be rendered for, or ``None``
when the Dag is not team-owned.
"""
link_teams = _get_extra_link_class_teams().get(type(link))
# Not registered by any plugin (defined by the operator), or registered by at least
# one global plugin: either way it is not restricted to a team.
if link_teams is None or None in link_teams:
return True
return team_name in link_teams
@cache
def get_timetables_plugins() -> dict[str, type[Timetable]]:
"""Collect and get timetable classes registered by plugins."""
@@ -17,18 +17,29 @@
from __future__ import annotations
import pytest
from sqlalchemy import update
from airflow._shared.timezones import timezone
from airflow.api_fastapi.common.dagbag import dag_bag_from_app
from airflow.api_fastapi.core_api.datamodels.extra_links import ExtraLinkCollectionResponse
from airflow.models.dag import DagModel
from airflow.models.dagbag import DBDagBag
from airflow.models.dagbundle import DagBundleModel
from airflow.models.team import Team
from airflow.models.xcom import XComModel as XCom
from airflow.plugins_manager import AirflowPlugin
from airflow.utils.state import DagRunState
from airflow.utils.types import DagRunTriggeredByType, DagRunType
from tests_common.test_utils.compat import BaseOperatorLink
from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_db_xcom
from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import (
clear_db_dag_bundles,
clear_db_dags,
clear_db_runs,
clear_db_teams,
clear_db_xcom,
)
from tests_common.test_utils.mock_operators import CustomOperator
pytestmark = pytest.mark.db_test
@@ -72,6 +83,46 @@ class TryNumberPlugin(AirflowPlugin):
operator_extra_links = [TryNumberLink()]
class TeamALink(BaseOperatorLink):
name = "Team A"
def get_link(self, operator, ti_key):
return "https://team-a.example.com"
class EveryoneLink(BaseOperatorLink):
name = "Everyone"
def get_link(self, operator, ti_key):
return "https://example.com/everyone"
class ImpersonatingLink(BaseOperatorLink):
"""Shares a name with ``CustomOperator``'s own ``CustomOpLink``."""
name = "Google Custom"
def get_link(self, operator, ti_key):
return "https://team-a.example.com/impersonated"
class TeamAPlugin(AirflowPlugin):
name = "team_a_plugin"
team_name = "team-a"
global_operator_extra_links = [TeamALink()]
class GlobalLinkPlugin(AirflowPlugin):
name = "global_link_plugin"
global_operator_extra_links = [EveryoneLink()]
class TeamANameCollisionPlugin(AirflowPlugin):
name = "team_a_name_collision_plugin"
team_name = "team-a"
global_operator_extra_links = [ImpersonatingLink()]
@pytest.mark.mock_plugin_manager(plugins=[])
class TestGetExtraLinks:
dag_id = "TEST_DAG_ID"
@@ -410,3 +461,111 @@ class TestGetExtraLinks:
params={"try_number": 99999},
)
assert response.status_code == 404
class TestGetExtraLinksTeamFiltering:
dag_id = "TEST_TEAM_LINKS_DAG"
dag_run_id = "TEST_TEAM_LINKS_RUN"
task_id = "TEST_TEAM_LINKS_TASK"
bundle_name = "team-links-bundle"
default_time = timezone.datetime(2020, 1, 1)
@staticmethod
def _clear_db():
clear_db_dags()
clear_db_runs()
clear_db_xcom()
clear_db_dag_bundles()
clear_db_teams()
@pytest.fixture(autouse=True)
def setup(self, test_client, dag_maker, session) -> None:
self._clear_db()
with dag_maker(
dag_id=self.dag_id,
schedule=None,
default_args={"start_date": self.default_time},
serialized=True,
):
CustomOperator(task_id=self.task_id, bash_command="TEST_LINK_VALUE")
dag_maker.create_dagrun(
run_id=self.dag_run_id,
logical_date=self.default_time,
run_type=DagRunType.MANUAL,
state=DagRunState.SUCCESS,
data_interval=(timezone.datetime(2020, 1, 1), timezone.datetime(2020, 1, 2)),
run_after=timezone.datetime(2020, 1, 2),
triggered_by=DagRunTriggeredByType.TEST,
)
test_client.app.dependency_overrides[dag_bag_from_app] = lambda: DBDagBag()
def teardown_method(self) -> None:
self._clear_db()
def _assign_dag_to_team(self, session, team_name: str | None) -> None:
"""Point the Dag's bundle at a bundle owned by ``team_name``, or by no team."""
bundle = DagBundleModel(name=self.bundle_name)
if team_name is not None:
bundle.teams.append(Team(name=team_name))
session.add(bundle)
session.flush()
session.execute(
update(DagModel).where(DagModel.dag_id == self.dag_id).values(bundle_name=self.bundle_name)
)
session.commit()
def _get_link_names(self, test_client) -> set[str]:
response = test_client.get(
f"/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_id}/links",
)
assert response.status_code == 200
return set(response.json()["extra_links"])
@pytest.mark.parametrize(
("multi_team", "dag_team", "expected_links"),
[
pytest.param(
"True",
"team-a",
{"Google Custom", "Everyone", "Team A"},
id="owning-team-sees-its-own-link",
),
pytest.param("True", "team-b", {"Google Custom", "Everyone"}, id="other-team-does-not-see-it"),
pytest.param("True", None, {"Google Custom", "Everyone"}, id="teamless-dag-does-not-see-it"),
pytest.param(
"False",
"team-b",
{"Google Custom", "Everyone", "Team A"},
id="unfiltered-when-multi-team-off",
),
],
)
@pytest.mark.mock_plugin_manager(plugins=[TeamAPlugin, GlobalLinkPlugin])
def test_team_scoped_links_filtered_by_dag_team(
self, test_client, session, multi_team, dag_team, expected_links
):
"""``Google Custom`` is the operator's own link and ``Everyone`` a global plugin's, so both
stay visible throughout; only ``Team A`` is team-scoped."""
with conf_vars({("core", "multi_team"): multi_team}):
self._assign_dag_to_team(session, dag_team)
assert self._get_link_names(test_client) == expected_links
@pytest.mark.parametrize("dag_team", ["team-a", "team-b"])
@pytest.mark.mock_plugin_manager(plugins=[TeamANameCollisionPlugin])
def test_operator_link_survives_a_team_plugin_reusing_its_name(self, test_client, session, dag_team):
"""A team plugin registering an existing link *name* must not make the operator's own link
team-scoped, which is why ownership is tracked per link class rather than per name."""
with conf_vars({("core", "multi_team"): "True"}):
self._assign_dag_to_team(session, dag_team)
response = test_client.get(
f"/dags/{self.dag_id}/dagRuns/{self.dag_run_id}/taskInstances/{self.task_id}/links",
)
assert response.status_code == 200
extra_links = response.json()["extra_links"]
assert set(extra_links) == {"Google Custom"}
assert extra_links["Google Custom"] != "https://team-a.example.com/impersonated"
@@ -912,3 +912,85 @@ class TestWarnAboutUnknownTranslationKeys:
plugins_manager.warn_about_unknown_translation_keys(plugin_translations, tmp_path / "en")
assert any("'a'" in record.getMessage() for record in caplog.records)
class TestExtraLinkTeamVisibility:
"""``is_extra_link_visible_to_team`` decides whether a team-scoped plugin's operator link
is rendered for a given Dag, so the API server can hide one team's links from another."""
@staticmethod
def _link_class():
from tests_common.test_utils.compat import BaseOperatorLink
class SomeLink(BaseOperatorLink):
name = "Some Link"
def get_link(self, operator, ti_key):
return "https://example.com"
return SomeLink
def test_operator_defined_link_is_visible_to_every_team(self):
"""A link no plugin registered belongs to the operator, so no team owns it."""
from airflow import plugins_manager
link = self._link_class()()
with mock_plugin_manager(plugins=[]):
assert plugins_manager.is_extra_link_visible_to_team(link, "team_a") is True
assert plugins_manager.is_extra_link_visible_to_team(link, None) is True
@pytest.mark.parametrize(
("dag_team", "expected"),
[
pytest.param("team_a", True, id="owning-team"),
pytest.param("team_b", False, id="other-team"),
pytest.param(None, False, id="teamless-dag"),
],
)
def test_team_scoped_link_is_visible_only_to_its_team(self, dag_team, expected):
from airflow import plugins_manager
link_class = self._link_class()
class TeamPlugin(AirflowPlugin):
name = "team_a_link_plugin"
team_name = "team_a"
global_operator_extra_links = [link_class()]
with mock_plugin_manager(plugins=[TeamPlugin]):
assert plugins_manager.is_extra_link_visible_to_team(link_class(), dag_team) is expected
def test_link_registered_by_a_global_plugin_too_stays_global(self):
"""Ownership resolves least restrictively: one global registration keeps the link
visible everywhere, rather than the team registration narrowing it."""
from airflow import plugins_manager
link_class = self._link_class()
class TeamPlugin(AirflowPlugin):
name = "team_a_link_plugin"
team_name = "team_a"
global_operator_extra_links = [link_class()]
class GlobalPlugin(AirflowPlugin):
name = "global_link_plugin"
global_operator_extra_links = [link_class()]
with mock_plugin_manager(plugins=[TeamPlugin, GlobalPlugin]):
assert plugins_manager.is_extra_link_visible_to_team(link_class(), "team_b") is True
assert plugins_manager.is_extra_link_visible_to_team(link_class(), None) is True
def test_operator_scoped_links_are_tracked_alongside_global_ones(self):
"""``operator_extra_links`` are team-owned on the same terms as ``global_operator_extra_links``."""
from airflow import plugins_manager
link_class = self._link_class()
class TeamPlugin(AirflowPlugin):
name = "team_a_link_plugin"
team_name = "team_a"
operator_extra_links = [link_class()]
with mock_plugin_manager(plugins=[TeamPlugin]):
assert plugins_manager.is_extra_link_visible_to_team(link_class(), "team_a") is True
assert plugins_manager.is_extra_link_visible_to_team(link_class(), "team_b") is False
@@ -93,6 +93,7 @@ def mock_plugin_manager(plugins=None, **kwargs):
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()