amoghrajesh commented on code in PR #47339:
URL: https://github.com/apache/airflow/pull/47339#discussion_r1998073615
##########
airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -224,16 +225,39 @@ def set_xcom(
# (the mapped task would fail in a moment as it can't be expanded anyway.)
# We use `BaseXCom.set` to set XComs directly to the database, bypassing
the XCom Backend.
+ from airflow.models.dagrun import DagRun
+
+ if not run_id:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"Run with ID:
`{run_id}` was not found")
+
+ dag_run_id = session.query(DagRun.id).filter_by(dag_id=dag_id,
run_id=run_id).scalar()
+ if dag_run_id is None:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG run not found on
DAG {dag_id} with ID {run_id}")
+
+ # Remove duplicate XComs and insert a new one.
+ session.execute(
+ delete(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,
+ )
+ )
+
try:
- BaseXCom.set(
+ # We expect serialised value from the caller - sdk, do not serialise
in here
+ new = cast(Any, XComModel)( # Work around Mypy complaining model not
defining '__init__'.
Review Comment:
Correct! We do not need it anymore as there's no magic here.
--
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]
