This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bd8cad6b729 Let non task callers set and delete variables over
execution API (#76375)
bd8cad6b729 is described below
commit bd8cad6b729c1680cbd1ebf271000a2bb9c8e746
Author: Amogh Desai <[email protected]>
AuthorDate: Wed Nov 4 16:43:56 2026 +0530
Let non task callers set and delete variables over execution API (#76375)
* Let non task callers set and delete variables over execution API
* review from jed
---
.../api_fastapi/execution_api/routes/variables.py | 3 +
.../airflow/api_fastapi/execution_api/security.py | 11 +++
.../execution_api/versions/v2026_10_30.py | 2 +-
.../api_fastapi/execution_api/test_security.py | 98 ++++++++++++----------
4 files changed, 68 insertions(+), 46 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 7f44b1dfc19..5a0f20beda2 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
@@ -34,6 +34,7 @@ from airflow.api_fastapi.execution_api.security import (
ExecutionAPIRoute,
ExecutionOrDagParseToken,
ExecutionOrProcessorSecretsToken,
+ allow_non_ti_writer,
get_team_name_dep,
)
from airflow.models.variable import Variable
@@ -147,6 +148,7 @@ def get_variable(
status.HTTP_403_FORBIDDEN: {"description": "The caller does not have
access to the variable"},
},
)
+@allow_non_ti_writer
def put_variable(
variable_key: Annotated[str, Path(min_length=1)],
body: VariablePostBody,
@@ -167,6 +169,7 @@ def put_variable(
status.HTTP_403_FORBIDDEN: {"description": "The caller 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 ff6daa52c6a..c76308b6398 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/security.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/security.py
@@ -145,6 +145,7 @@ _REQUEST_SCOPE_JOB_KEY = "dag_processor_job_id"
_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")
@@ -157,6 +158,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.
@@ -288,6 +295,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
@@ -477,12 +485,14 @@ class ExecutionAPIRoute(APIRoute):
allowed_token_types: frozenset[str]
skip_auto_ti_attempt_live: bool
+ allow_non_ti_writer: bool
requires_open_job: 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:
@@ -537,6 +547,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 1f7780ee321..1dbebc7fb41 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 6839530c6b5..77c2b604366 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
@@ -1104,18 +1104,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),
@@ -1127,8 +1158,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"),
]
@@ -1146,8 +1175,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
@@ -1155,7 +1184,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
@@ -1163,10 +1191,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()
@@ -1190,15 +1219,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(
@@ -1273,14 +1299,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
@@ -1307,19 +1331,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
@@ -1348,9 +1362,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"
@@ -1365,10 +1377,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
