kaxil commented on code in PR #59874:
URL: https://github.com/apache/airflow/pull/59874#discussion_r3707446313
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -572,3 +574,42 @@ def delete(
map_index=map_index,
),
)
+
+ @classmethod
+ def delete_all(
+ cls,
+ dag_id: str,
+ run_id: str,
+ task_id: str | None = None,
+ key: str | None = None,
+ map_index: int | None = None,
+ ) -> int:
+ """
+ Bulk delete XCom entries, optionally filtered by task_id, key, or
map_index.
+
+ :param dag_id: Dag ID.
+ :param run_id: Dag run ID for the task.
+ :param task_id: Optional task ID filter. If provided, only XComs from
this task
+ will be deleted. Pass *None* (default) to delete across all tasks.
+ :param key: Optional key filter. If provided, only XComs with this key
+ will be deleted. Pass *None* (default) to delete all keys.
+ :param map_index: Optional map index filter. If provided, only XComs
with this
+ map index will be deleted. Pass *None* (default) to delete all map
indexes.
+ :return: Number of XCom entries deleted.
+ """
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ msg = SUPERVISOR_COMMS.send(
Review Comment:
`delete` calls `cls.purge(xcom_result)` before sending the message, which is
how a custom backend drops the payload it wrote outside the database
(`XComObjectStorageBackend.purge` unlinks the object storage file).
`delete_all` skips that, so on a custom backend the rows go away and the
objects behind them are orphaned with nothing left pointing at them, which is
the population this feature is aimed at. Can it purge the matched entries, or
at minimum carry the note `XComModel.clear` has: "This **will not** purge any
data from a custom XCom backend"? The docs sentence that would need amending is
xcoms.rst:113, "This will be called as part of ``delete``".
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -572,3 +574,42 @@ def delete(
map_index=map_index,
),
)
+
+ @classmethod
+ def delete_all(
+ cls,
+ dag_id: str,
+ run_id: str,
+ task_id: str | None = None,
+ key: str | None = None,
+ map_index: int | None = None,
+ ) -> int:
+ """
+ Bulk delete XCom entries, optionally filtered by task_id, key, or
map_index.
+
+ :param dag_id: Dag ID.
+ :param run_id: Dag run ID for the task.
+ :param task_id: Optional task ID filter. If provided, only XComs from
this task
+ will be deleted. Pass *None* (default) to delete across all tasks.
+ :param key: Optional key filter. If provided, only XComs with this key
+ will be deleted. Pass *None* (default) to delete all keys.
+ :param map_index: Optional map index filter. If provided, only XComs
with this
+ map index will be deleted. Pass *None* (default) to delete all map
indexes.
+ :return: Number of XCom entries deleted.
Review Comment:
Worth a caveat in here: with no `task_id` or `key` filter this also removes
the `dag_result=True` rows that `GET /dags/{dag_id}/dagRuns/{dag_run_id}/wait`
serves. `_serialize_response` guards with `if result_xcoms := await
self._serialize_xcoms():`, so a client waiting on the run gets `{"state":
"success"}` with the `results` key absent rather than an empty dict, and the
end-of-dag cleanup task described in the PR body is exactly the shape that
triggers it. Skipping `dag_result` rows unless they are asked for, or at least
documenting it and steering people to the filters, would avoid a silent hole
that is hard to trace back here.
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -392,6 +394,7 @@ def from_api_response(cls, response: HITLDetailResponse) ->
HITLDetailResponseRe
| GetVariableKeys
| PutVariable
| DeleteXCom
+ | BulkDeleteXCom
Review Comment:
Adding `BulkDeleteXCom` here makes it sendable from a trigger, but
`XComDeleteCountResult` never went into `ToTriggerRunner` above and the test
parks it in `in_task_but_not_in_trigger_runner` instead. Since
`handle_bulk_delete_xcom` returns a response (unlike `handle_delete_xcom`,
which returns `None`, which is why `DeleteXCom` needs no result type) and the
runner validates replies with `TypeAdapter(ToTriggerRunner)`, a trigger calling
`delete_all` gets `Input tag 'XComDeleteCountResult' does not match any of the
expected tags` after the rows are already gone. The one comparable pair is
`GetPreviousTI`/`PreviousTIResult`, but `get_previous_ti` only exists on
`RuntimeTaskInstance`, so it has no trigger-side caller. `delete_all` is a
classmethod callable from anywhere, and `TriggerRunner.init_comms` sets
`task_runner.SUPERVISOR_COMMS = self.comms_decoder`, so the call does go
through from a trigger. Adding `| XComDeleteCountResult` to `ToTriggerRunner`
and dropping the test exclusio
n fixes it, or the triggerer wiring could come back out.
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -652,6 +653,27 @@ def delete(
# decouple from the server response string
return OKResponse(ok=True)
+ def delete_all(
+ self,
+ dag_id: str,
+ run_id: str,
+ task_id: str | None = None,
+ key: str | None = None,
+ map_index: int | None = None,
+ ) -> XComDeleteCountResult:
+ """Bulk delete XCom values via the API server."""
+ params: dict[str, str | int] = {}
+
+ if map_index is not None:
+ params["map_index"] = map_index
+ if task_id is not None:
+ params["task_id"] = task_id
+ if key is not None:
+ params["key"] = key
+
+ resp = self.client.delete(url=f"xcoms/{dag_id}/{run_id}",
params=params)
+ return XComDeleteCountResult(count=resp.json()["count"])
Review Comment:
`Client.request` retries on `httpx.RequestError` and 5xx, and the server
commits inside `create_session` before the response is written, so a DELETE
whose response is lost gets replayed and the retry finds nothing left to
delete. The caller then sees `count=0` after N rows actually went away, which
`delete` above avoids only because it returns a fixed `OKResponse`. Worth
documenting the count as best effort rather than exact, given the alternative
is an idempotency key?
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -482,3 +482,36 @@ def delete_xcom(
)
session.execute(query)
return {"message": f"XCom with key: {key} successfully deleted."}
+
+
[email protected](
+ "/{dag_id}/{run_id}",
+ description="Bulk delete Xcom values.",
+)
+def bulk_delete_xcoms(
+ session: SessionDep,
+ dag_id: str,
+ run_id: str,
+ task_id: Annotated[str | None, Query()] = None,
+ key: Annotated[str | None, Query()] = None,
+ map_index: Annotated[int | None, Query()] = None,
+):
+ """Bulk delete Xcom values."""
+ query = delete(XComModel).where(
+ XComModel.dag_id == dag_id,
+ XComModel.run_id == run_id,
+ )
+
+ if task_id is not None:
+ query = query.where(XComModel.task_id == task_id)
+
+ if key is not None:
+ query = query.where(XComModel.key == key)
+
+ if map_index is not None:
+ query = query.where(XComModel.map_index == map_index)
+
+ result = session.execute(query)
+ count = getattr(result, "rowcount", 0)
Review Comment:
Two things about producing and returning this count. `session.execute()` on
a DELETE returns a `CursorResult`, which always carries `rowcount`, so the
`getattr` fallback cannot fire, and if it ever could then reporting 0 would
tell the caller nothing was deleted while rows were. Returning a bare dict also
keeps the shape out of the OpenAPI spec, so codegen emits no model and the
client has to reach for `resp.json()["count"]`; the sibling count routes in
this API (`get_dr_count`, `get_task_instance_count`) just annotate `-> int` and
let the client wrap it, as `TICount(count=resp.json())` does.
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_09_09.py:
##########
@@ -0,0 +1,30 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from cadwyn import VersionChange, endpoint
+
+
+class AddXcomBulkDeleteEndpoint(VersionChange):
+ """Add XCom bulk delete endpoint."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ endpoint("/xcoms/{dag_id}/{run_id}", ["DELETE"]).didnt_exist,
Review Comment:
Nothing asserts the 404 contract this version file promises.
`versions/v2026_06_30/test_variables.py` has the pattern for the equivalent new
endpoint change: pin `Airflow-API-Version` to the previous release and assert
404. The versioning guide asks for tests covering both the new and the previous
version, and this is the only thing that pins what older workers see.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -482,3 +482,36 @@ def delete_xcom(
)
session.execute(query)
return {"message": f"XCom with key: {key} successfully deleted."}
+
+
[email protected](
+ "/{dag_id}/{run_id}",
+ description="Bulk delete Xcom values.",
+)
+def bulk_delete_xcoms(
+ session: SessionDep,
+ dag_id: str,
+ run_id: str,
+ task_id: Annotated[str | None, Query()] = None,
+ key: Annotated[str | None, Query()] = None,
+ map_index: Annotated[int | None, Query()] = None,
+):
+ """Bulk delete Xcom values."""
+ query = delete(XComModel).where(
Review Comment:
`XComModel.clear` uses the same denormalized predicate so the shape is
established, but the delta here is that without `task_id` the seek degrades to
the `dag_id` prefix of `idx_xcom_task_instance(dag_id, task_id, run_id,
map_index)`, so this walks the dag's whole XCom history to delete one run's
worth and holds the locks for that long. `XComModel.set` already resolves the
run first with `select(DagRun.id).where(DagRun.dag_id == ..., DagRun.run_id ==
...)`, and `dag_run_id` is the leading primary key column, so filtering on that
instead would make it an exact range. Once per dag run is not a hot path, but
the deployments this feature targets are the ones with the biggest xcom tables.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py:
##########
@@ -617,6 +625,111 @@ def test_xcom_delete_endpoint(self, client,
create_task_instance, session):
).first()
assert xcom_ti is not None
+ @pytest.mark.parametrize(
+ ("task_id", "key", "expected_remaining", "expected_deleted"),
+ [
+ pytest.param(None, None, 0, 4, id="all_xcoms_for_run"),
+ pytest.param("t1", None, 2, 2, id="all_keys_for_task"),
+ pytest.param(None, "xcom_3", 3, 1, id="specific_key_all_tasks"),
+ ],
+ )
+ def test_xcom_bulk_delete_endpoint(
+ self, client, dag_maker, session, task_id, key, expected_remaining,
expected_deleted
+ ):
+ """Test XCom bulk deletions."""
+
+ with dag_maker(dag_id="dag"):
+ EmptyOperator(task_id="t1")
+ EmptyOperator(task_id="t2")
+
+ dag_run = dag_maker.create_dagrun(run_id="test")
+
+ ti = dag_run.get_task_instance("t1")
+ ti2 = dag_run.get_task_instance("t2")
+
+ ti.xcom_push(key="xcom_1", value='"value1"', session=session)
+ ti.xcom_push(key="xcom_2", value='"value2"', session=session)
+
+ ti2.xcom_push(key="xcom_1", value='"value1"', session=session)
+ ti2.xcom_push(key="xcom_3", value='"value3"', session=session)
+ session.commit()
+
+ params = {}
+ if task_id is not None:
+ params["task_id"] = task_id
+ if key is not None:
+ params["key"] = key
+ response = client.delete(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}",
params=params)
+
+ assert response.status_code == 200
+ assert response.json() == {"count": expected_deleted}
+
+ xcoms = session.scalars(
+ select(XComModel).where(XComModel.dag_id == ti.dag_id,
XComModel.run_id == ti.run_id)
+ ).all()
+ assert len(xcoms) == expected_remaining
+
+ if task_id == "t1" and key is None:
+ assert not any(xcom.task_id == "t1" for xcom in xcoms)
+ assert all(xcom.task_id == "t2" for xcom in xcoms)
+
+ remaining_keys = {xcom.key for xcom in xcoms}
+ assert remaining_keys == {"xcom_1", "xcom_3"}
+
+ elif task_id is None and key == "xcom_3":
+ assert not any(xcom.key == "xcom_3" for xcom in xcoms)
+ assert all(xcom.key != "xcom_3" for xcom in xcoms)
+
+ remaining_tasks = {xcom.task_id for xcom in xcoms}
+ assert remaining_tasks == {"t1", "t2"}
+
+ remaining_keys = {xcom.key for xcom in xcoms}
+ assert remaining_keys == {"xcom_1", "xcom_2"}
+
+ def test_xcom_bulk_delete_by_map_index(self, client, dag_maker, session):
+ """Test XCom bulk deletion by map_index."""
+
+ class MyOperator(EmptyOperator):
+ def __init__(self, *, x, **kwargs):
+ super().__init__(**kwargs)
+ self.x = x
+
+ with dag_maker(dag_id="dag"):
+ MyOperator.partial(task_id="t1").expand(x=[1, 2])
+ MyOperator.partial(task_id="t2").expand(x=[1])
+
+ dag_run = dag_maker.create_dagrun(run_id="test")
+ tis = {(ti.task_id, ti.map_index): ti for ti in dag_run.task_instances}
+
+ for task_id, map_index in (("t1", 0), ("t1", 1), ("t2", 0)):
+ ti = tis[(task_id, map_index)]
+ session.add(
+ XComModel(
+ key="xcom_1",
+ value='"value1"',
+ dag_run_id=ti.dag_run.id,
+ run_id=ti.run_id,
+ task_id=ti.task_id,
+ dag_id=ti.dag_id,
+ map_index=map_index,
+ )
+ )
+ session.commit()
+
+ response = client.delete(
+ f"/execution/xcoms/{dag_run.dag_id}/{dag_run.run_id}",
params={"map_index": 0}
+ )
+
+ assert response.status_code == 200
+ assert response.json() == {"count": 2}
+
+ remaining = session.scalars(
+ select(XComModel.map_index).where(
+ XComModel.dag_id == dag_run.dag_id, XComModel.run_id ==
dag_run.run_id
+ )
+ ).all()
+ assert set(remaining) == {1}
+
class TestXComTeamAccess:
Review Comment:
Two gaps in the new tests. Both bulk cases create a single dag run, so
removing `XComModel.run_id == run_id` from the route's WHERE clause would leave
every assertion green; `test_xcom_delete_endpoint` just above does the right
thing by creating a second dag and asserting its XCom survives. And
`TestXComTeamAccess._url` below only builds the single XCom path, so the new
write route never passes through `has_xcom_access` in any test even though this
PR changes that function's signature.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
