This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch changes-3.4.0rc3 in repository https://gitbox.apache.org/repos/asf/airflow.git
commit aa149990e7a680a31be80b4c81f8310e120a1798 Author: Rahul Vats <[email protected]> AuthorDate: Wed Nov 4 21:00:10 2026 +0530 Let non task callers set and delete variables over execution API (#76375) (#76436) * Let non task callers set and delete variables over execution API * review from jed (cherry picked from commit bd8cad6b729c1680cbd1ebf271000a2bb9c8e746) Co-authored-by: Amogh Desai <[email protected]> --- .../api_fastapi/execution_api/routes/variables.py | 11 ++- .../airflow/api_fastapi/execution_api/security.py | 11 +++ .../execution_api/versions/v2026_10_30.py | 2 +- .../api_fastapi/execution_api/test_security.py | 104 ++++++++++++--------- 4 files changed, 80 insertions(+), 48 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py index 5b1c15cdd27..9964c4fe754 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py @@ -29,7 +29,12 @@ from airflow.api_fastapi.execution_api.datamodels.variable import ( VariablePostBody, VariableResponse, ) -from airflow.api_fastapi.execution_api.security import CurrentTIToken, get_team_name_dep +from airflow.api_fastapi.execution_api.security import ( + CurrentTIToken, + ExecutionAPIRoute, + allow_non_ti_writer, + get_team_name_dep, +) from airflow.models.variable import Variable @@ -57,7 +62,7 @@ async def has_variable_access( return True -router = APIRouter() +router = APIRouter(route_class=ExecutionAPIRoute) log = logging.getLogger(__name__) @@ -139,6 +144,7 @@ def get_variable( status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the variable"}, }, ) +@allow_non_ti_writer def put_variable( variable_key: Annotated[str, Path(min_length=1)], body: VariablePostBody, @@ -159,6 +165,7 @@ def put_variable( status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the variable"}, }, ) +@allow_non_ti_writer def delete_variable( variable_key: Annotated[str, Path(min_length=1)], team_name: Annotated[str | None, Depends(get_team_name_dep)], diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/security.py b/airflow-core/src/airflow/api_fastapi/execution_api/security.py index f209fe2da02..09bb22f2910 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py @@ -99,6 +99,7 @@ _REQUEST_SCOPE_LIVE_ATTEMPT_KEY = "live_attempt_checked" _IN_PROCESS_NON_TI_CALLER = "airflow_in_process_non_ti_caller" _IN_PROCESS_TRIGGERER_ID = "airflow_in_process_triggerer_id" _SKIP_AUTO_TI_ATTEMPT_LIVE = "skip_auto_ti_attempt_live" +_ALLOW_NON_TI_WRITER = "allow_non_ti_writer" _P = ParamSpec("_P") _R = TypeVar("_R") @@ -109,6 +110,12 @@ def skip_auto_ti_attempt_live(endpoint: Callable[_P, _R]) -> Callable[_P, _R]: return endpoint +def allow_non_ti_writer(endpoint: Callable[_P, _R]) -> Callable[_P, _R]: + """Mark a route whose data no task attempt owns, so ``require_auth`` skips its automatic check.""" + setattr(endpoint, _ALLOW_NON_TI_WRITER, True) + return endpoint + + class JWTBearer(HTTPBearer): """ Validates JWT tokens for the Execution API. @@ -232,6 +239,7 @@ async def require_auth( and not request.scope.get(_IN_PROCESS_NON_TI_CALLER) and not request.scope.get(_REQUEST_SCOPE_LIVE_ATTEMPT_KEY) and not getattr(route, _SKIP_AUTO_TI_ATTEMPT_LIVE, False) + and not getattr(route, _ALLOW_NON_TI_WRITER, False) ): # The versions package imports routes, which depend on this module. from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates @@ -304,10 +312,12 @@ class ExecutionAPIRoute(APIRoute): allowed_token_types: frozenset[str] skip_auto_ti_attempt_live: bool + allow_non_ti_writer: bool def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.skip_auto_ti_attempt_live = bool(getattr(self.endpoint, _SKIP_AUTO_TI_ATTEMPT_LIVE, False)) + self.allow_non_ti_writer = bool(getattr(self.endpoint, _ALLOW_NON_TI_WRITER, False)) all_scopes: set[str] = set() for dep in self.dependencies: @@ -361,6 +371,7 @@ def _team_name_for_ti_stmt(ti_id): .join(DagBundleModel, DagBundleModel.name == DagModel.bundle_name) .join(DagBundleModel.teams) .where(TaskInstance.id == ti_id) + .execution_options(include_all_attempts=True) ) 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 22b34e5115d..7c6942aa8ef 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 @@ -72,7 +72,7 @@ class AddStoppedTaskReport(VersionChange): class IdentifyArchivedTaskStateUpdates(VersionChangeWithSideEffects): - """Reject every mutation from an archived attempt with 410; older clients keep each endpoint's own response.""" + """Reject mutations from an archived attempt with 410, except Variable writes; older clients keep each endpoint's own response.""" description = __doc__ instructions_to_migrate_to_previous_version = () diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py index 39b01143093..7d24a1cd446 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_security.py @@ -41,10 +41,13 @@ from airflow.api_fastapi.execution_api.security import ( get_team_name_dep, require_auth, ) +from airflow.models import DagModel from airflow.models.callback import Callback, CallbackFetchMethod from airflow.models.connection_test import ConnectionTestRequest, ConnectionTestState +from airflow.models.dagbundle import DagBundleModel from airflow.models.renderedtifields import RenderedTaskInstanceFields from airflow.models.taskinstance import LegacyTaskDataOwner, TaskInstance +from airflow.models.team import Team from airflow.models.variable import Variable from airflow.models.xcom import XComModel, XComModelV1, XComModelV2 from airflow.sdk.api.client import XComOperations @@ -53,6 +56,9 @@ from airflow.sdk.execution_time import request_handlers, task_runner from airflow.sdk.execution_time.comms import DeleteXCom, GetXCom from airflow.utils.state import CallbackState, TaskInstanceState +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import clear_db_teams + class TestTIClaims: def test_defaults_scope_and_retains_extra(self): @@ -352,18 +358,49 @@ class TestAttemptLiveness: assert client.delete(path).status_code == 410 assert client.get(path).status_code == 410 - def test_unknown_execution_identity_cannot_write(self, client, caller): - _, token = caller + def test_unknown_execution_identity_cannot_write_xcom(self, client, caller, session): + ti, token = caller token.id = uuid4() + path = f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/result" + + assert client.post(path, json="unowned").status_code == 404 + assert session.scalar(select(XComModel).where(XComModel.dag_id == ti.dag_id)) is None + + @pytest.mark.parametrize("identity", ["unknown", "archived"]) + def test_non_ti_identity_can_write_variables(self, client, caller, session, identity): + ti, token = caller + if identity == "archived": + ti.prepare_db_for_next_try(session) + session.commit() + else: + token.id = uuid4() key = f"attempt-liveness-{uuid4()}" try: - response = client.put(f"/execution/variables/{key}", json={"value": "unowned"}) - assert response.status_code == 404 + assert client.put(f"/execution/variables/{key}", json={"value": "unowned"}).status_code == 201 + assert Variable.get(key) == "unowned" + assert client.delete(f"/execution/variables/{key}").status_code == 204 with pytest.raises(KeyError): Variable.get(key) finally: Variable.delete(key) + def test_archived_attempt_writes_variable_to_its_team(self, client, caller, session): + ti, _ = caller + bundle_name = session.scalar(select(DagModel.bundle_name).where(DagModel.dag_id == ti.dag_id)) + bundle = session.get(DagBundleModel, bundle_name) + bundle.teams.append(Team(name="team_a")) + ti.prepare_db_for_next_try(session) + session.commit() + key = f"attempt-liveness-{uuid4()}" + try: + with conf_vars({("core", "multi_team"): "True"}): + response = client.put(f"/execution/variables/{key}", json={"value": "zombie"}) + assert response.status_code == 201 + assert session.scalar(select(Variable.team_name).where(Variable.key == key)) == "team_a" + finally: + Variable.delete(key) + clear_db_teams() + ARCHIVED_ROUTES = [ ("put", "/execution/store/ti/{ti_id}/key", {"value": "stale"}), ("delete", "/execution/store/ti/{ti_id}/key", None), @@ -375,8 +412,6 @@ class TestAttemptLiveness: ("delete", "/execution/store/asset/by-name/clear?name=asset", None), ("post", "/execution/dag-runs/target/run", {}), ("post", "/execution/dag-runs/target/run/clear", None), - ("put", "/execution/variables/key", {"value": "stale"}), - ("delete", "/execution/variables/key", None), ("post", "/execution/xcoms/{dag_id}/{run_id}/{task_id}/key", "late"), ] @@ -394,8 +429,8 @@ class TestAttemptLiveness: assert response.status_code == 410, response.text - def test_callback_execution_token_can_mutate_variable(self, client, caller, session): - _, token = caller + def test_callback_execution_token_can_write_xcom(self, client, caller, session): + ti, token = caller callback = Callback() callback.fetch_method = CallbackFetchMethod.IMPORT_PATH callback.state = CallbackState.QUEUED @@ -403,7 +438,6 @@ class TestAttemptLiveness: session.commit() token.id = callback.id token.claims.scope = "callback" - key = f"attempt-liveness-{uuid4()}" try: response = client.patch(f"/execution/callbacks/{callback.id}/run") assert response.status_code == 204 @@ -411,10 +445,11 @@ class TestAttemptLiveness: token.id = UUID(claims["sub"]) token.claims = TIClaims(**claims) - assert client.put(f"/execution/variables/{key}", json={"value": "callback"}).status_code == 201 - assert Variable.get(key) == "callback" + response = client.post( + f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/key", json="callback" + ) + assert response.status_code == 201, response.text finally: - Variable.delete(key) session.delete(callback) session.commit() @@ -438,15 +473,12 @@ class TestAttemptLiveness: session.delete(connection_test) session.commit() - def test_trusted_in_process_caller_without_ti_can_write(self): - key = f"attempt-liveness-{uuid4()}" - try: - with TestClient(InProcessExecutionAPI().app) as client: - response = client.put(f"/variables/{key}", json={"value": "watcher"}) - assert response.status_code == 201 - assert Variable.get(key) == "watcher" - finally: - Variable.delete(key) + def test_trusted_in_process_caller_without_ti_can_write(self, caller, session): + ti, _ = caller + with TestClient(InProcessExecutionAPI().app) as client: + response = client.post(f"/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/key", json="watcher") + assert response.status_code == 201, response.text + assert XComModelV2.get_for_attempt(ti.id, "key", session=session).value == "watcher" @pytest.mark.parametrize("delete_rejected", [False, True], ids=["accepted", "archived-token-rejected"]) def test_trusted_in_process_xcom_uses_current_owner_without_attempt_token( @@ -521,14 +553,12 @@ class TestAttemptLiveness: ) assert XComModelV2.get_for_attempt(archived.id, "key", session=session).value == "archived" - @pytest.mark.parametrize("operation", ["xcom_write", "xcom_delete", "rtif_write", "variable_write"]) + @pytest.mark.parametrize("operation", ["xcom_write", "xcom_delete", "rtif_write"]) def test_admitted_child_mutation_stays_with_archived_uuid( - self, client, caller, session, monkeypatch, request, operation + self, client, caller, session, monkeypatch, operation ): attempt, _ = caller old_id = attempt.id - variable_key = f"attempt-liveness-{uuid4()}" - request.addfinalizer(lambda: Variable.delete(variable_key)) xcom_path = f"/execution/xcoms/{attempt.dag_id}/{attempt.run_id}/{attempt.task_id}/value" XComModel.set_for_attempt( task_instance_id=old_id, key="value", value="original", serialize=False, session=session @@ -555,19 +585,9 @@ class TestAttemptLiveness: ) ) session.commit() - path = { - "rtif_write": f"/execution/task-instances/{old_id}/rtif", - "variable_write": f"/execution/variables/{variable_key}", - }.get(operation, xcom_path) - method = { - "xcom_write": "POST", - "xcom_delete": "DELETE", - "rtif_write": "PUT", - "variable_write": "PUT", - }[operation] - body = {"rtif_write": {"field": "late"}, "variable_write": {"value": "in flight"}}.get( - operation, "late" - ) + path = {"rtif_write": f"/execution/task-instances/{old_id}/rtif"}.get(operation, xcom_path) + method = {"xcom_write": "POST", "xcom_delete": "DELETE", "rtif_write": "PUT"}[operation] + body = {"rtif_write": {"field": "late"}}.get(operation, "late") admitted, finish = Event(), Event() check = security._require_live_attempt @@ -596,9 +616,7 @@ class TestAttemptLiveness: finally: finish.set() response = pending.result(timeout=10) - expected_status = {"xcom_write": 201, "xcom_delete": 200, "rtif_write": 410, "variable_write": 201}[ - operation - ] + expected_status = {"xcom_write": 201, "xcom_delete": 200, "rtif_write": 410}[operation] assert response.status_code == expected_status, response.text session.expire_all() assert XComModelV2.get_for_attempt(successor.id, "value", session=session).value == "successor" @@ -613,10 +631,6 @@ class TestAttemptLiveness: session.get(XComModelV1, (attempt.dag_run.id, attempt.task_id, attempt.map_index, "value")) is None ) - elif operation == "variable_write": - assert Variable.get(variable_key) == "in flight" - assert client.put(path, json={"value": "too late"}).status_code == 410 - assert Variable.get(variable_key) == "in flight" else: assert RenderedTaskInstanceFields.get_for_attempt(old_id, session=session) is None
