justinpakzad commented on code in PR #59874:
URL: https://github.com/apache/airflow/pull/59874#discussion_r3724669913
##########
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:
I have added an another test that creates a second dag run and asserts it's
XCom survives. I also added two additional tests to `TestXComTeamAccess` for
the bulk path.
##########
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:
Good find. I've updated it to first resolve the run and then to filter on
`dag_run_id` as suggested.
--
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]
