This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 2cbd7b0b8ec Let non task callers set and delete variables over
execution API (#76375) (#76436)
2cbd7b0b8ec is described below
commit 2cbd7b0b8ece8ffe8d005a17d0a8a59720f1f5a0
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
