mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 07:12:48 +00:00
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:
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()
|
||||
|
||||
Reference in new issue
Block a user