From 63b615b252589ef49b8504f2dbe2942b543ec896 Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Tue, 22 Sep 2026 09:43:16 -0700 Subject: [PATCH] Scope plugin macros to their team in multi-team mode (#73224) A team-scoped plugin's macros were reachable from every task's templates, so any team could call macros written to talk to another team's systems. The worker decides this from a new ``multi_team`` flag on the run context rather than from its own config, which is not guaranteed to match the scheduler's (misconfigured or intentional attempt to read macros from other teams): reading it as disabled while it is in fact enabled would drop scoping and apply a team's macros everywhere. The existing ``team_name`` cannot carry that signal, being ``None`` both for a teamless task and for every task when multi-team is off, so a plugin declaring a team would otherwise lose its macros in single-team deployments. Scoping is applied at attribute access rather than when macros are loaded, because the loader mutates a process-wide module and a worker process can serve tasks from more than one team. --- .../execution_api/datamodels/taskinstance.py | 11 +++ .../execution_api/routes/task_instances.py | 1 + .../execution_api/versions/__init__.py | 8 +- .../execution_api/versions/v2026_10_30.py | 13 ++++ .../versions/head/test_task_instances.py | 43 +++++++++++ .../v2025_04_28/test_task_instances.py | 2 + .../tests_common/test_utils/mock_plugins.py | 1 + .../airflow/sdk/api/datamodels/_generated.py | 1 + .../src/airflow/sdk/execution_time/context.py | 21 ++++++ .../sdk/execution_time/schema/schema.json | 5 ++ .../airflow/sdk/execution_time/task_runner.py | 6 +- task-sdk/src/airflow/sdk/plugins_manager.py | 13 ++++ .../task_sdk/execution_time/test_context.py | 74 +++++++++++++++++++ .../execution_time/test_task_runner.py | 33 +++++++++ ts-sdk/src/generated/supervisor.ts | 2 + 15 files changed, 232 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ddf31db9718..a0f0f5e3553 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -447,6 +447,17 @@ class TIRunContext(BaseModel): ``None`` for regular tasks and for stub tasks that declare no parameters. """ + multi_team: bool = False + """ + Whether the deployment runs in multi-team mode. + + Sent explicitly because a worker cannot read ``core.multi_team`` itself: its config is + not guaranteed to match the scheduler's, and reading it as disabled while it is in fact + enabled would drop team scoping and apply a team's plugins to every task. ``team_name`` + cannot stand in for this, being ``None`` both for a teamless task and for every task + when multi-team is off. + """ + class PrevSuccessfulDagRunResponse(BaseModel): """Schema for response with previous successful DagRun information for Task Template Context.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 14910a4d866..3dd4bc11a0b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -314,6 +314,7 @@ def ti_run( connections=[], xcom_keys_to_clear=xcom_keys, should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), + multi_team=conf.getboolean("core", "multi_team"), ) # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 70ea9be2be6..79f3f1833fc 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -55,11 +55,17 @@ from airflow.api_fastapi.execution_api.versions.v2026_06_30 import ( from airflow.api_fastapi.execution_api.versions.v2026_10_30 import ( AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, + AddMultiTeamToTIRunContext, ) bundle = VersionBundle( HeadVersion(), - Version("2026-10-30", AddArgBindingsToTIRunContext, AddCallbackRunEndpoint), + Version( + "2026-10-30", + AddArgBindingsToTIRunContext, + AddCallbackRunEndpoint, + AddMultiTeamToTIRunContext, + ), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index 0620053b3e2..6d5730a71ab 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -52,3 +52,16 @@ class AddCallbackRunEndpoint(VersionChange): instructions_to_migrate_to_previous_version = ( endpoint("/callbacks/{callback_id}/run", ["PATCH"]).didnt_exist, ) + + +class AddMultiTeamToTIRunContext(VersionChange): + """Add ``multi_team`` so a worker can determine multi-team (e.g. for plugin scoping) without needing to trust its own config.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("multi_team").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_multi_team_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Strip ``multi_team`` from the run context for older clients.""" + response.body.pop("multi_team", None) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 80278c28d07..8ecf7b1b5dd 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -316,6 +316,7 @@ class TestTIRunState: "variables": [], "connections": [], "xcom_keys_to_clear": [], + "multi_team": False, } # upstream_map_indexes is now computed by Task SDK, not returned by the server in HEAD version assert "upstream_map_indexes" not in result @@ -831,6 +832,7 @@ class TestTIRunState: "variables": [], "connections": [], "xcom_keys_to_clear": [], + "multi_team": False, "next_method": "execute_complete", "next_kwargs": expected_next_kwargs, "start_date": None, @@ -906,6 +908,7 @@ class TestTIRunState: "variables": [], "connections": [], "xcom_keys_to_clear": [], + "multi_team": False, "next_method": "execute_complete", "next_kwargs": expected_next_kwargs, } @@ -1116,6 +1119,46 @@ class TestTIRunState: assert dag_run["run_id"] == "test" assert dag_run["state"] == "running" + @pytest.mark.parametrize( + ("multi_team_enabled", "expected"), + [ + pytest.param("False", False, id="multi-team-disabled"), + pytest.param("True", True, id="multi-team-enabled"), + ], + ) + def test_ti_run_reports_multi_team( + self, client, session, create_task_instance, time_machine, multi_team_enabled, expected + ): + """The worker cannot read ``core.multi_team`` itself, so the run context carries it.""" + instant_str = "2024-09-30T12:00:00Z" + instant = timezone.parse(instant_str) + time_machine.move_to(instant, tick=False) + + ti = create_task_instance( + task_id="test_ti_run_reports_multi_team", + state=State.QUEUED, + dagrun_state=DagRunState.RUNNING, + session=session, + start_date=instant, + dag_id=str(uuid4()), + ) + session.commit() + + with conf_vars({("core", "multi_team"): multi_team_enabled}): + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": instant_str, + }, + ) + + assert response.status_code == 200 + assert response.json()["multi_team"] is expected + @pytest.mark.parametrize( ("multi_team_enabled", "expect_team"), [ diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py index efa29338d92..40954978cc8 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py @@ -96,3 +96,5 @@ class TestTIUpdateState: assert result["task_reschedule_count"] == 0 assert result["max_tries"] == 0 assert result["should_retry"] is False + # Added in 2026-10-30; older clients must not see it. + assert "multi_team" not in result 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 fb3c674dbf0..cc100747dd0 100644 --- a/devel-common/src/tests_common/test_utils/mock_plugins.py +++ b/devel-common/src/tests_common/test_utils/mock_plugins.py @@ -99,6 +99,7 @@ def mock_plugin_manager(plugins=None, **kwargs): 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() if plugins is not None or "import_errors" in kwargs: diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index de4b2ebe420..a435be7e165 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -829,3 +829,4 @@ class TIRunContext(BaseModel): should_retry: Annotated[bool | None, Field(title="Should Retry")] = False start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None arg_bindings: Annotated[list[TaskArgBinding] | None, Field(title="Arg Bindings")] = None + multi_team: Annotated[bool | None, Field(title="Multi Team")] = False diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 5a14238b088..fdb696ec906 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -995,6 +995,14 @@ class MacrosAccessor: """Wrapper to access Macros module lazily.""" _macros_module = None + # Class-level defaults so a plain ``MacrosAccessor()`` keeps working and attribute + # lookup never falls through to ``__getattr__`` (which would recurse). + _team_name: str | None = None + _multi_team: bool = False + + def __init__(self, team_name: str | None = None, multi_team: bool = False) -> None: + self._team_name = team_name + self._multi_team = multi_team def __getattr__(self, item: str) -> Any: # Lazily load Macros module @@ -1002,6 +1010,19 @@ class MacrosAccessor: import airflow.sdk.execution_time.macros self._macros_module = airflow.sdk.execution_time.macros + + if self._multi_team: + from airflow.sdk.plugins_manager import get_macro_plugin_teams + + owning_team = get_macro_plugin_teams().get(item) + # ``None`` covers both a global plugin and an attribute that is not a plugin + # module at all, such as a built-in macro. + if owning_team is not None and owning_team != self._team_name: + raise AttributeError( + f"Macros of plugin {item!r} belong to team {owning_team!r} and are not " + f"available to this task." + ) + return getattr(self._macros_module, item) def __repr__(self) -> str: diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index dc6a8c053ce..7ae8b5db6e6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -5121,6 +5121,11 @@ ], "default": null, "title": "Arg Bindings" + }, + "multi_team": { + "default": false, + "title": "Multi Team", + "type": "boolean" } }, "required": [ diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 8460d6bbb5e..cacd6ae87ab 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -303,8 +303,12 @@ class RuntimeTaskInstance(TaskInstance): integrate_macros_plugins() dag_run_conf: dict[str, Any] | None = None + macros_accessor = MacrosAccessor() if from_server := self._ti_context_from_server: dag_run_conf = from_server.dag_run.conf or dag_run_conf + macros_accessor = MacrosAccessor( + team_name=from_server.dag_run.team_name, multi_team=bool(from_server.multi_team) + ) validated_params = process_params(self.task.dag, self.task, dag_run_conf, suppress_exception=False) @@ -323,7 +327,7 @@ class RuntimeTaskInstance(TaskInstance): "ti": self, "outlet_events": OutletEventAccessors(), "inlet_events": InletEventsAccessors(self.task.inlets), - "macros": MacrosAccessor(), + "macros": macros_accessor, "params": validated_params, # TODO: Make this go through Public API longer term. # "test_mode": task_instance.test_mode, diff --git a/task-sdk/src/airflow/sdk/plugins_manager.py b/task-sdk/src/airflow/sdk/plugins_manager.py index 3934e54ef5e..f5acbe60858 100644 --- a/task-sdk/src/airflow/sdk/plugins_manager.py +++ b/task-sdk/src/airflow/sdk/plugins_manager.py @@ -132,6 +132,19 @@ def integrate_macros_plugins() -> None: ) +@cache +def get_macro_plugin_teams() -> dict[str, str | None]: + """ + Map the name of each plugin contributing macros to the team owning it. + + Macros are attached to one module per plugin name, so this is what lets a task be + offered its own team's and the global plugins' macros but not another team's. Only + plugins that actually contribute macros get a module, hence a submodule to hide. + """ + plugins, _ = _get_plugins() + return {plugin.name: plugin.team_name for plugin in plugins if plugin.name and plugin.macros} + + def integrate_listener_plugins(listener_manager: ListenerManager) -> None: """Add listeners from plugins.""" plugins, _ = _get_plugins() diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index cbc579c1f15..3c8c5924423 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -86,6 +86,7 @@ from airflow.sdk.execution_time.context import ( AssetStateStoreAccessors, ConnectionAccessor, InletEventsAccessors, + MacrosAccessor, OutletEventAccessor, OutletEventAccessors, TaskStateStoreAccessor, @@ -104,6 +105,7 @@ from airflow.sdk.execution_time.secrets import ExecutionAPISecretsBackend from airflow.sdk.state import BaseStoreBackend from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.mock_plugins import mock_plugin_manager if TYPE_CHECKING: from pydantic import JsonValue @@ -2538,3 +2540,75 @@ class TestAssetStateStoreAccessorWithCustomBackend: result = await AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark") assert result == "2026-05-01" + + +class TestMacrosAccessorTeamScoping: + """A task may use its own team's and the global plugins' macros, but not another team's.""" + + @staticmethod + def _plugins(): + from airflow.sdk.plugins_manager import AirflowPlugin + + def team_a_macro(): + return "team-a" + + def shared_macro(): + return "shared" + + class TeamAPlugin(AirflowPlugin): + name = "team_a_macros" + team_name = "team-a" + macros = [team_a_macro] + + class GlobalPlugin(AirflowPlugin): + name = "global_macros" + macros = [shared_macro] + + return [TeamAPlugin, GlobalPlugin] + + @pytest.fixture + def integrated_macros(self): + from airflow.sdk.plugins_manager import integrate_macros_plugins + + with mock_plugin_manager(plugins=self._plugins()): + integrate_macros_plugins() + yield + + @pytest.mark.parametrize( + ("team_name", "reachable"), + [ + pytest.param("team-a", True, id="owning-team"), + pytest.param("team-b", False, id="other-team"), + pytest.param(None, False, id="teamless-task"), + ], + ) + def test_team_macros_reachable_only_by_their_team(self, integrated_macros, team_name, reachable): + accessor = MacrosAccessor(team_name=team_name, multi_team=True) + + if reachable: + assert accessor.team_a_macros.team_a_macro() == "team-a" + else: + with pytest.raises(AttributeError, match="belong to team 'team-a'"): + accessor.team_a_macros + + @pytest.mark.parametrize( + "team_name", + [pytest.param("team-a", id="team-task"), pytest.param(None, id="teamless-task")], + ) + def test_global_plugin_macros_stay_reachable(self, integrated_macros, team_name): + accessor = MacrosAccessor(team_name=team_name, multi_team=True) + + assert accessor.global_macros.shared_macro() == "shared" + + def test_builtin_macros_stay_reachable(self, integrated_macros): + """Only plugin submodules are scoped; the macros module's own contents are not.""" + accessor = MacrosAccessor(team_name="team-b", multi_team=True) + + assert accessor.ds_add("2026-01-01", 1) == "2026-01-02" + + def test_nothing_is_hidden_when_multi_team_is_off(self, integrated_macros): + """A plugin declaring a team in a single-team deployment keeps working as before.""" + accessor = MacrosAccessor() + + assert accessor.team_a_macros.team_a_macro() == "team-a" + assert accessor.global_macros.shared_macro() == "shared" diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 17ba0709ca0..faceabb5c72 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2272,6 +2272,39 @@ class TestRuntimeTaskInstance: "ti": runtime_ti, } + def test_macros_in_context_are_scoped_to_the_tasks_team(self, create_runtime_ti, mock_supervisor_comms): + """The accessor placed in the context must carry the team the server reported.""" + from airflow.sdk.plugins_manager import AirflowPlugin + + from tests_common.test_utils.mock_plugins import mock_plugin_manager + + def team_a_macro(): + return "team-a" + + class TeamAPlugin(AirflowPlugin): + name = "team_a_macros" + team_name = "team-a" + macros = [team_a_macro] + + runtime_ti = create_runtime_ti(task=BaseOperator(task_id="hello"), dag_id="basic_task") + # Stand in for a multi-team server handing this task to a worker as team-b's. + runtime_ti._ti_context_from_server.multi_team = True + runtime_ti._ti_context_from_server.dag_run.team_name = "team-b" + + dr = runtime_ti._ti_context_from_server.dag_run + mock_supervisor_comms.send.return_value = PrevSuccessfulDagRunResult( + data_interval_end=dr.logical_date - timedelta(hours=1), + data_interval_start=dr.logical_date - timedelta(hours=2), + start_date=dr.start_date - timedelta(hours=1), + end_date=dr.start_date, + ) + + with mock_plugin_manager(plugins=[TeamAPlugin]): + macros = runtime_ti.get_template_context()["macros"] + + with pytest.raises(AttributeError, match="belong to team 'team-a'"): + macros.team_a_macros + def test_get_context_with_ti_context_from_server(self, create_runtime_ti, mock_supervisor_comms): """Test the context keys are added when sent from API server (mocked)""" diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 9b89cbf804c..263dfa4c50b 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -257,6 +257,7 @@ export type TaskId1 = string; export type Name9 = string; export type Kind1 = "literal"; export type FromDefault = boolean; +export type MultiTeam = boolean; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -1032,6 +1033,7 @@ export interface TIRunContext { should_retry?: ShouldRetry; start_date?: StartDate2; arg_bindings?: ArgBindings; + multi_team?: MultiTeam; } /** * Variable schema for responses with fields that are needed for Runtime.