Handle concurrent XCom writes with savepoint and UPDATE fallback to avoid 409

airflow-core/newsfragments/69959.bugfix.rst

@@ -0,0 +1 @@
+Fix concurrent or retried XCom writes returning HTTP 409 by using a savepoint and UPDATE fallback.

airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py

@@ -22,7 +22,8 @@
 
 from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status
 from pydantic import JsonValue
-from sqlalchemy import delete
+from sqlalchemy import delete, update
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.sql.selectable import Select
 
 from airflow.api_fastapi.common.db.common import SessionDep
@@ -433,17 +434,50 @@ def set_xcom(
     # means loading the serialized dag and that seems like a relatively costly operation for minimal benefit
     # (the mapped task would fail in a moment as it can't be expanded anyway.)
     try:
-        # We expect serialised value from the caller - sdk, do not serialise in here
-        XComModel.set(
-            key=key,
-            value=value,
-            run_id=run_id,
-            task_id=task_id,
-            dag_id=dag_id,
-            map_index=map_index,
-            serialize=False,
-            dag_result=dag_result,
-            session=session,
+        # Use a savepoint so that an IntegrityError on the XCom write does not
+        # roll back the task-map merge that may have already been flushed above.
+        with session.begin_nested():
+            # We expect serialised value from the caller - sdk, do not serialise in here
+            XComModel.set(
+                key=key,
+                value=value,
+                run_id=run_id,
+                task_id=task_id,
+                dag_id=dag_id,
+                map_index=map_index,
+                serialize=False,
+                dag_result=dag_result,
+                session=session,
+            )
+    except IntegrityError as e:
+        # Only swallow unique-constraint violations (a concurrent or retried write
+        # for the same PK beat us to it). FK violations — e.g. the task_instance
+        # cascade-deleted between our dag_run_id lookup and the flush — must
+        # propagate; silently returning 200 there would hide a lost write.
+        if not any(
+            msg in str(e.orig)
+            for msg in (
+                "UNIQUE constraint failed",  # SQLite
+                "Duplicate entry",  # MySQL
+                "violates unique constraint",  # PostgreSQL
+            )
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=f"Unexpected database error: {e.orig}",
+            )
+        # The savepoint was rolled back automatically; fall through to an explicit
+        # UPDATE so the latest value wins.
+        session.execute(
+            update(XComModel)
+            .where(
+                XComModel.key == key,
+                XComModel.run_id == run_id,
+                XComModel.task_id == task_id,
+                XComModel.dag_id == dag_id,
+                XComModel.map_index == map_index,
+            )
+            .values(value=value, dag_result=dag_result)
         )
     except ValueError as e:
         raise HTTPException(status.HTTP_404_NOT_FOUND, str(e))

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py

@@ -579,6 +579,89 @@ def test_xcom_dag_result(self, client, create_task_instance, session):
         )
         assert dag_result is True
 
+    def test_xcom_set_concurrent_write_uses_update_fallback(self, client, create_task_instance, session):
+        """
+        When a concurrent write wins the PK race and XComModel.set raises a
+        unique-constraint IntegrityError, the savepoint is rolled back and the
+        endpoint falls through to an UPDATE. The caller should still get a 201 and
+        the row should hold the value from this (latest) write.
+        """
+        from unittest.mock import patch
+
+        from sqlalchemy.exc import IntegrityError
+
+        ti = create_task_instance()
+        session.commit()
+
+        # Pre-seed the row so the UPDATE fallback has a target row to update.
+        session.add(
+            XComModel(
+                key="concurrent_key",
+                value='"original"',
+                dag_run_id=ti.dag_run.id,
+                run_id=ti.run_id,
+                task_id=ti.task_id,
+                dag_id=ti.dag_id,
+            )
+        )
+        session.commit()
+
+        # Simulate the concurrent-write race: the first INSERT hits a unique
+        # constraint because a parallel write already committed the same PK.
+        unique_orig = Exception("violates unique constraint (xcom_pkey)")
+        unique_error = IntegrityError("INSERT INTO xcom ...", {}, unique_orig)
+
+        with patch(
+            "airflow.api_fastapi.execution_api.routes.xcoms.XComModel.set",
+            side_effect=unique_error,
+        ):
+            response = client.post(
+                f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/concurrent_key",
+                json='"updated"',
+            )
+
+        assert response.status_code == 201
+        assert response.json() == {"message": "XCom successfully set"}
+
+        stored = session.scalar(
+            select(XComModel.value).where(
+                XComModel.dag_id == ti.dag_id,
+                XComModel.run_id == ti.run_id,
+                XComModel.task_id == ti.task_id,
+                XComModel.key == "concurrent_key",
+            )
+        )
+        assert stored == '"updated"', f"Expected '\"updated\"' but got {stored!r}"
+
+    def test_xcom_set_fk_integrity_error_propagates(self, client, create_task_instance, session):
+        """
+        A non-unique IntegrityError — e.g. an FK violation when the task_instance
+        cascade-deleted between the dag_run_id lookup and the flush — must NOT be
+        swallowed. Silently returning 201 would hide the fact that the XCom was
+        never actually stored.
+        """
+        from unittest.mock import patch
+
+        from sqlalchemy.exc import IntegrityError
+
+        ti = create_task_instance()
+        session.commit()
+
+        fk_orig = Exception("FOREIGN KEY constraint failed")
+        fk_error = IntegrityError("INSERT INTO xcom ...", {}, fk_orig)
+
+        with patch(
+            "airflow.api_fastapi.execution_api.routes.xcoms.XComModel.set",
+            side_effect=fk_error,
+        ):
+            response = client.post(
+                f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/fk_key",
+                json='"value"',
+            )
+
+        # FK errors must not be silenced — they should propagate as a server error.
+        assert response.status_code == 500
+
 
 class TestXComsDeleteEndpoint:
     def test_xcom_delete_endpoint(self, client, create_task_instance, session):