This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new f624c4e94d2 [v3-3-test] Fix task state store Execution API rejecting
keys that contain slashes (#69178) (#70967)
f624c4e94d2 is described below
commit f624c4e94d2011ed5fb8cfee040331e213106c00
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Tue Aug 4 11:09:14 2026 +0530
[v3-3-test] Fix task state store Execution API rejecting keys that contain
slashes (#69178) (#70967)
* Fix task state store API key path
* Revert uv.lock
* Add client side
* Simplified the test
(cherry picked from commit 4899e9041859db58fabeb87f231e454a5f65a418)
Co-authored-by: Takayoshi Makabe
<[email protected]>
---
.../execution_api/routes/task_state_store.py | 6 ++--
.../execution_api/versions/v2026_06_30.py | 6 ++--
.../versions/head/test_task_state_store.py | 26 ++++++++++++++
task-sdk/src/airflow/sdk/api/client.py | 6 ++--
task-sdk/tests/task_sdk/api/test_client.py | 41 ++++++++++++++++++++++
5 files changed, 76 insertions(+), 9 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
index 624926fe639..e1dfe7fcd97 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_state_store.py
@@ -58,7 +58,7 @@ def _get_task_scope_for_ti(task_instance_id: UUID, session:
Session) -> TaskScop
return TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id,
map_index=ti.map_index)
[email protected]("/{task_instance_id}/{key}")
[email protected]("/{task_instance_id}/{key:path}")
def get_task_state_store(
task_instance_id: UUID,
key: Annotated[str, Path(min_length=1)],
@@ -78,7 +78,7 @@ def get_task_state_store(
return TaskStateStoreResponse(value=json.loads(value))
[email protected]("/{task_instance_id}/{key}",
status_code=status.HTTP_204_NO_CONTENT)
[email protected]("/{task_instance_id}/{key:path}",
status_code=status.HTTP_204_NO_CONTENT)
def set_task_state_store(
task_instance_id: UUID,
key: Annotated[str, Path(min_length=1)],
@@ -90,7 +90,7 @@ def set_task_state_store(
get_state_backend().set(scope, key, json.dumps(body.value),
expires_at=body.expires_at, session=session)
[email protected]("/{task_instance_id}/{key}",
status_code=status.HTTP_204_NO_CONTENT)
[email protected]("/{task_instance_id}/{key:path}",
status_code=status.HTTP_204_NO_CONTENT)
def delete_task_state_store(
task_instance_id: UUID,
key: Annotated[str, Path(min_length=1)],
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py
index e89e2ed04cc..930bd1426a4 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_06_30.py
@@ -114,9 +114,9 @@ class AddTaskAndAssetStateStoreEndpoints(VersionChange):
description = __doc__
instructions_to_migrate_to_previous_version = (
- endpoint("/store/ti/{task_instance_id}/{key}", ["GET"]).didnt_exist,
- endpoint("/store/ti/{task_instance_id}/{key}", ["PUT"]).didnt_exist,
- endpoint("/store/ti/{task_instance_id}/{key}", ["DELETE"]).didnt_exist,
+ endpoint("/store/ti/{task_instance_id}/{key:path}",
["GET"]).didnt_exist,
+ endpoint("/store/ti/{task_instance_id}/{key:path}",
["PUT"]).didnt_exist,
+ endpoint("/store/ti/{task_instance_id}/{key:path}",
["DELETE"]).didnt_exist,
endpoint("/store/ti/{task_instance_id}", ["DELETE"]).didnt_exist,
endpoint("/store/asset/by-name/value", ["GET"]).didnt_exist,
endpoint("/store/asset/by-name/value", ["PUT"]).didnt_exist,
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
index f544de4f95e..671ceffbb4a 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_state_store.py
@@ -78,6 +78,15 @@ class TestGetTaskState:
assert response.status_code == 404
assert "Task instance" in response.json()["detail"]["message"]
+ def test_get_key_with_slash(self, client: TestClient,
create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "spark/job_id"), json={"value":
"spark_001"})
+
+ response = client.get(_api_url(ti.id, "spark/job_id"))
+
+ assert response.status_code == 200
+ assert response.json() == {"value": "spark_001"}
+
class TestPutTaskState:
def test_put_creates_row(self, client: TestClient, create_task_instance:
CreateTaskInstance):
@@ -203,6 +212,14 @@ class TestPutTaskState:
assert response.status_code == 404
+ def test_put_key_with_slash(self, client: TestClient,
create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+
+ response = client.put(_api_url(ti.id, "spark/job_id"), json={"value":
"spark_001"})
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "spark/job_id")).json() == {"value":
"spark_001"}
+
class TestDeleteTaskState:
def test_delete_removes_key(self, client: TestClient,
create_task_instance: CreateTaskInstance):
@@ -231,6 +248,15 @@ class TestDeleteTaskState:
assert client.get(_api_url(ti.id, "job_id")).status_code == 404
assert client.get(_api_url(ti.id, "checkpoint")).json() == {"value":
"b"}
+ def test_delete_key_with_slash(self, client: TestClient,
create_task_instance: CreateTaskInstance):
+ ti = create_task_instance()
+ client.put(_api_url(ti.id, "spark/job_id"), json={"value":
"spark_001"})
+
+ response = client.delete(_api_url(ti.id, "spark/job_id"))
+
+ assert response.status_code == 204
+ assert client.get(_api_url(ti.id, "spark/job_id")).status_code == 404
+
class TestClearTaskState:
def test_clear_removes_all_keys(self, client: TestClient,
create_task_instance: CreateTaskInstance):
diff --git a/task-sdk/src/airflow/sdk/api/client.py
b/task-sdk/src/airflow/sdk/api/client.py
index 037ddd7eb5f..8e0a46a477a 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -724,7 +724,7 @@ class TaskStateStoreOperations:
def get(self, ti_id: uuid.UUID, key: str) -> TaskStateStoreResponse |
ErrorResponse:
"""Get a task store value from the API server."""
try:
- resp = self.client.get(f"store/ti/{ti_id}/{key}")
+ resp = self.client.get(f"store/ti/{ti_id}/{quote(key, safe='')}")
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.debug("Task store key not found", ti_id=ti_id, key=key)
@@ -735,12 +735,12 @@ class TaskStateStoreOperations:
def set(self, ti_id: uuid.UUID, key: str, value: JsonValue, expires_at:
datetime | None) -> OKResponse:
"""Set a task store value via the API server."""
body = TaskStateStorePutBody(value=value, expires_at=expires_at)
- self.client.put(f"store/ti/{ti_id}/{key}",
content=body.model_dump_json())
+ self.client.put(f"store/ti/{ti_id}/{quote(key, safe='')}",
content=body.model_dump_json())
return OKResponse(ok=True)
def delete(self, ti_id: uuid.UUID, key: str) -> OKResponse:
"""Delete a single task store key via the API server."""
- self.client.delete(f"store/ti/{ti_id}/{key}")
+ self.client.delete(f"store/ti/{ti_id}/{quote(key, safe='')}")
return OKResponse(ok=True)
def clear(self, ti_id: uuid.UUID) -> OKResponse:
diff --git a/task-sdk/tests/task_sdk/api/test_client.py
b/task-sdk/tests/task_sdk/api/test_client.py
index 92b970f7ee5..5c4bac56972 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -1831,6 +1831,19 @@ class TestTaskStateOperations:
assert isinstance(result, TaskStateStoreResponse)
assert result.value == "spark_app_001"
+ def test_get_url_encodes_key_with_slash(self):
+ requests_seen = []
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ requests_seen.append(request)
+ return httpx.Response(status_code=200, json={"value":
"spark_app_001"})
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ client.task_state_store.get(ti_id=self.TI_ID, key="spark/job_id")
+
+ assert b"%2F" in requests_seen[0].url.raw_path
+ assert requests_seen[0].url.raw_path ==
f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
+
def test_get_returns_error_response_on_404(self):
def handle_request(request: httpx.Request) -> httpx.Response:
return httpx.Response(
@@ -1858,6 +1871,21 @@ class TestTaskStateOperations:
)
assert result == OKResponse(ok=True)
+ def test_set_url_encodes_key_with_slash(self):
+ requests_seen = []
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ requests_seen.append(request)
+ return httpx.Response(status_code=204)
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ client.task_state_store.set(
+ ti_id=self.TI_ID, key="spark/job_id", value="spark_app_001",
expires_at=None
+ )
+
+ assert b"%2F" in requests_seen[0].url.raw_path
+ assert requests_seen[0].url.raw_path ==
f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
+
def test_set_with_expires_at_sends_field(self):
"""expires_at is forwarded as an ISO datetime string in the request
body."""
expires = datetime(2026, 5, 21, 12, 0, 0, tzinfo=dt_timezone.utc)
@@ -1901,6 +1929,19 @@ class TestTaskStateOperations:
result = client.task_state_store.delete(ti_id=self.TI_ID, key="job_id")
assert result == OKResponse(ok=True)
+ def test_delete_url_encodes_key_with_slash(self):
+ requests_seen = []
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ requests_seen.append(request)
+ return httpx.Response(status_code=204)
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ client.task_state_store.delete(ti_id=self.TI_ID, key="spark/job_id")
+
+ assert b"%2F" in requests_seen[0].url.raw_path
+ assert requests_seen[0].url.raw_path ==
f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
+
def test_clear_sends_delete_request(self):
def handle_request(request: httpx.Request) -> httpx.Response:
assert request.method == "DELETE"
