Re: [PR] fix: handle concurrent XCom writes with savepoint + UPDATE fallback to avoid 409 [airflow]

["nagasrisai (via GitHub)" <gi...@apache.org>]

nagasrisai commented on code in PR #69959:
URL: https://github.com/apache/airflow/pull/69959#discussion_r3713955961


##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -433,17 +434,35 @@ 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:

Review Comment:
   Really good catch, thank you. You're right — a FK violation on the task_instance cascade would also produce an `IntegrityError` here, and swallowing it silently would return 200 when the XCom was never actually stored. I'll narrow the guard to unique-constraint violations only, mirroring `_UniqueConstraintErrorHandler`, and let everything else propagate. Will push the fix shortly.



##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -433,17 +434,35 @@ 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():

Review Comment:
   Agreed, the fallback path needs a direct test. I'll add one to `test_xcoms.py` that pre-seeds the row and patches `XComModel.set` to raise `IntegrityError` (unique constraint) once, then verifies the endpoint returns 200 and the stored value is the latest one. That way the UPDATE fallback is actually exercised rather than assumed.



-- 
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: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org