AIP-72: Handle Custom XCom Backend on Task SDK

airflow/api_fastapi/core_api/openapi/v1-generated.yaml

@@ -4641,7 +4641,7 @@ paths:
         required: false
         schema:
           type: boolean
-          default: true
+          default: false
           title: Stringify
       responses:
         '200':

airflow/api_fastapi/core_api/routes/public/xcom.py

@@ -37,7 +37,8 @@
 from airflow.api_fastapi.core_api.security import ReadableXComFilterDep, requires_access_dag
 from airflow.api_fastapi.logging.decorators import action_logging
 from airflow.exceptions import TaskNotFound
-from airflow.models import DAG, DagRun as DR, XCom
+from airflow.models import DAG, DagRun as DR
+from airflow.models.xcom import XComModel
 from airflow.settings import conf
 
 xcom_router = AirflowRouter(
@@ -63,22 +64,25 @@ def get_xcom_entry(
     session: SessionDep,
     map_index: Annotated[int, Query(ge=-1)] = -1,
     deserialize: Annotated[bool, Query()] = False,
-    stringify: Annotated[bool, Query()] = True,
+    stringify: Annotated[bool, Query()] = False,
 ) -> XComResponseNative | XComResponseString:
     """Get an XCom entry."""
     if deserialize:
         if not conf.getboolean("api", "enable_xcom_deserialize_support", fallback=False):
             raise HTTPException(
                 status.HTTP_400_BAD_REQUEST, "XCom deserialization is disabled in configuration."
             )
-        query = select(XCom, XCom.value)
+        query = select(XComModel, XComModel.value)
     else:
-        query = select(XCom)
+        query = select(XComModel)
 
     query = query.where(
-        XCom.dag_id == dag_id, XCom.task_id == task_id, XCom.key == xcom_key, XCom.map_index == map_index
+        XComModel.dag_id == dag_id,
+        XComModel.task_id == task_id,
+        XComModel.key == xcom_key,
+        XComModel.map_index == map_index,
     )
-    query = query.join(DR, and_(XCom.dag_id == DR.dag_id, XCom.run_id == DR.run_id))
+    query = query.join(DR, and_(XComModel.dag_id == DR.dag_id, XComModel.run_id == DR.run_id))
     query = query.where(DR.run_id == dag_run_id)
 
     if deserialize:
@@ -90,6 +94,8 @@ def get_xcom_entry(
         raise HTTPException(status.HTTP_404_NOT_FOUND, f"XCom entry with key: `{xcom_key}` not found")
 
     if deserialize:
+        from airflow.sdk.execution_time.xcom import XCom
+
         xcom, value = item
         xcom_stub = copy.copy(xcom)
         xcom_stub.value = value
@@ -127,19 +133,19 @@ def get_xcom_entries(
 
     This endpoint allows specifying `~` as the dag_id, dag_run_id, task_id to retrieve XCom entries for all DAGs.
     """
-    query = select(XCom)
+    query = select(XComModel)
     if dag_id != "~":
-        query = query.where(XCom.dag_id == dag_id)
-    query = query.join(DR, and_(XCom.dag_id == DR.dag_id, XCom.run_id == DR.run_id))
+        query = query.where(XComModel.dag_id == dag_id)
+    query = query.join(DR, and_(XComModel.dag_id == DR.dag_id, XComModel.run_id == DR.run_id))
 
     if task_id != "~":
-        query = query.where(XCom.task_id == task_id)
+        query = query.where(XComModel.task_id == task_id)
     if dag_run_id != "~":
         query = query.where(DR.run_id == dag_run_id)
     if map_index is not None:
-        query = query.where(XCom.map_index == map_index)
+        query = query.where(XComModel.map_index == map_index)
     if xcom_key is not None:
-        query = query.where(XCom.key == xcom_key)
+        query = query.where(XComModel.key == xcom_key)
 
     query, total_entries = paginated_select(
         statement=query,
@@ -148,7 +154,9 @@ def get_xcom_entries(
         limit=limit,
         session=session,
     )
-    query = query.order_by(XCom.dag_id, XCom.task_id, XCom.run_id, XCom.map_index, XCom.key)
+    query = query.order_by(
+        XComModel.dag_id, XComModel.task_id, XComModel.run_id, XComModel.map_index, XComModel.key
+    )
     xcoms = session.scalars(query)
     return XComCollectionResponse(xcom_entries=xcoms, total_entries=total_entries)
 
@@ -197,38 +205,48 @@ def create_xcom_entry(
         )
 
     # Check existing XCom
-    if XCom.get_one(
+    already_existing_query = XComModel.get_many(
         key=request_body.key,
-        task_id=task_id,
-        dag_id=dag_id,
+        task_ids=task_id,
+        dag_ids=dag_id,
         run_id=dag_run_id,
-        map_index=request_body.map_index,
+        map_indexes=request_body.map_index,
         session=session,
-    ):
+    )
+    result = already_existing_query.with_entities(XComModel.value).first()
+    if result:
         raise HTTPException(
             status_code=status.HTTP_409_CONFLICT,
             detail=f"The XCom with key: `{request_body.key}` with mentioned task instance already exists.",
         )
 
-    # Create XCom entry
-    XCom.set(
-        dag_id=dag_id,
-        task_id=task_id,
-        run_id=dag_run_id,
+    try:
+        value = XComModel.serialize_value(request_body.value)
+    except (ValueError, TypeError):
+        raise HTTPException(
+            status.HTTP_400_BAD_REQUEST, f"Couldn't serialise the XCom with key: `{request_body.key}`"
+        )
+
+    new = XComModel(
+        dag_run_id=dag_run.id,
         key=request_body.key,
-        value=XCom.serialize_value(request_body.value),
+        value=value,
+        run_id=dag_run_id,
+        task_id=task_id,
+        dag_id=dag_id,
         map_index=request_body.map_index,
-        session=session,
     )
+    session.add(new)
+    session.flush()
 
     xcom = session.scalar(
-        select(XCom)
+        select(XComModel)
         .filter(
-            XCom.dag_id == dag_id,
-            XCom.task_id == task_id,
-            XCom.run_id == dag_run_id,
-            XCom.key == request_body.key,
-            XCom.map_index == request_body.map_index,
+            XComModel.dag_id == dag_id,
+            XComModel.task_id == task_id,
+            XComModel.run_id == dag_run_id,
+            XComModel.key == request_body.key,
+            XComModel.map_index == request_body.map_index,
         )
         .limit(1)
     )
@@ -260,15 +278,15 @@ def update_xcom_entry(
 ) -> XComResponseNative:
     """Update an existing XCom entry."""
     # Check if XCom entry exists
-    xcom_new_value = XCom.serialize_value(patch_body.value)
+    xcom_new_value = XComModel.serialize_value(patch_body.value)
     xcom_entry = session.scalar(
-        select(XCom)
+        select(XComModel)
         .where(
-            XCom.dag_id == dag_id,
-            XCom.task_id == task_id,
-            XCom.run_id == dag_run_id,
-            XCom.key == xcom_key,
-            XCom.map_index == patch_body.map_index,
+            XComModel.dag_id == dag_id,
+            XComModel.task_id == task_id,
+            XComModel.run_id == dag_run_id,
+            XComModel.key == xcom_key,
+            XComModel.map_index == patch_body.map_index,
         )
         .limit(1)
     )
@@ -280,6 +298,6 @@ def update_xcom_entry(
         )
 
     # Update XCom entry
-    xcom_entry.value = XCom.serialize_value(xcom_new_value)
+    xcom_entry.value = XComModel.serialize_value(xcom_new_value)
 
     return XComResponseNative.model_validate(xcom_entry)

airflow/api_fastapi/core_api/security.py

@@ -43,7 +43,7 @@
 from airflow.models.dag import DagModel, DagRun, DagTag
 from airflow.models.dagwarning import DagWarning
 from airflow.models.taskinstance import TaskInstance as TI
-from airflow.models.xcom import XCom
+from airflow.models.xcom import XComModel
 
 if TYPE_CHECKING:
     from sqlalchemy.sql import Select
@@ -132,7 +132,7 @@ class PermittedXComFilter(PermittedDagFilter):
     """A parameter that filters the permitted XComs for the user."""
 
     def to_orm(self, select: Select) -> Select:
-        return select.where(XCom.dag_id.in_(self.value))
+        return select.where(XComModel.dag_id.in_(self.value))
 
 
 class PermittedTagFilter(PermittedDagFilter):

airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -273,6 +273,9 @@ class TIRunContext(BaseModel):
     Can either be a "decorated" dict, or a string encrypted with the shared Fernet key.
     """
 
+    xcom_keys_to_clear: Annotated[list[str], Field(default_factory=list)]
+    """List of Xcom keys that need to be cleared and purged on by the worker."""
+
 
 class PrevSuccessfulDagRunResponse(BaseModel):
     """Schema for response with previous successful DagRun information for Task Template Context."""

airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -48,7 +48,7 @@
 from airflow.models.taskinstance import TaskInstance as TI, _update_rtif
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
-from airflow.models.xcom import XCom
+from airflow.models.xcom import XComModel
 from airflow.utils import timezone
 from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
 
@@ -187,18 +187,22 @@ def ti_run(
         if not dr:
             raise ValueError(f"DagRun with dag_id={ti.dag_id} and run_id={ti.run_id} not found.")
 
-        # Clear XCom data for the task instance since we are certain it is executing
+        # Send the keys to the SDK so that the client requests to clear those XComs from the server.
+        # The reason we cannot do this here in the server is because we need to issue a purge on custom XCom backends
+        # too. With the current assumption, the workers ONLY have access to the custom XCom backends directly and they
+        # can issue the purge.
+
         # However, do not clear it for deferral
+        xcom_keys = []
         if not ti.next_method:
             map_index = None if ti.map_index < 0 else ti.map_index
-            log.info("Clearing xcom data for task id: %s", ti_id_str)
-            XCom.clear(
-                dag_id=ti.dag_id,
-                task_id=ti.task_id,
-                run_id=ti.run_id,
-                map_index=map_index,
-                session=session,
+            query = session.query(XComModel.key).filter_by(
+                dag_id=ti.dag_id, task_id=ti.task_id, run_id=ti.run_id
             )
+            if map_index is not None:
+                query = query.filter_by(map_index=map_index)
+
+            xcom_keys = [row.key for row in session.execute(query).all()]
 
         task_reschedule_count = (
             session.query(
@@ -216,6 +220,7 @@ def ti_run(
             # TODO: Add variables and connections that are needed (and has perms) for the task
             variables=[],
             connections=[],
+            xcom_keys_to_clear=xcom_keys,
         )
 
         # Only set if they are non-null

airflow/api_fastapi/execution_api/routes/xcoms.py

@@ -22,6 +22,7 @@
 
 from fastapi import Body, Depends, HTTPException, Query, Response, status
 from pydantic import JsonValue
+from sqlalchemy import delete
 from sqlalchemy.sql.selectable import Select
 
 from airflow.api_fastapi.common.db.common import SessionDep
@@ -30,7 +31,7 @@
 from airflow.api_fastapi.execution_api.datamodels.token import TIToken
 from airflow.api_fastapi.execution_api.datamodels.xcom import XComResponse
 from airflow.models.taskmap import TaskMap
-from airflow.models.xcom import BaseXCom
+from airflow.models.xcom import XComModel
 from airflow.utils.db import get_query_count
 
 # TODO: Add dependency on JWT token
@@ -62,7 +63,7 @@ async def xcom_query(
             },
         )
 
-    query = BaseXCom.get_many(
+    query = XComModel.get_many(
         run_id=run_id,
         key=key,
         task_ids=task_id,
@@ -126,7 +127,7 @@ def get_xcom(
     """Get an Airflow XCom from database - not other XCom Backends."""
     # The xcom_query allows no map_index to be passed. This endpoint should always return just a single item,
     # so we override that query value
-    xcom_query = xcom_query.filter(BaseXCom.map_index == map_index)
+    xcom_query = xcom_query.filter(XComModel.map_index == map_index)
     # We use `BaseXCom.get_many` to fetch XComs directly from the database, bypassing the XCom Backend.
     # This avoids deserialization via the backend (e.g., from a remote storage like S3) and instead
     # retrieves the raw serialized value from the database. By not relying on `XCom.get_many` or `XCom.get_one`
@@ -222,18 +223,39 @@ def set_xcom(
     # TODO: Can/should we check if a client _hasn't_ provided this for an upstream of a mapped task? That
     # 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.)
+    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,
+        )
+    )
 
-    # We use `BaseXCom.set` to set XComs directly to the database, bypassing the XCom Backend.
     try:
-        BaseXCom.set(
+        # We expect serialised value from the caller - sdk, do not serialise in here
+        new = XComModel(
+            dag_run_id=dag_run_id,
             key=key,
             value=value,
-            dag_id=dag_id,
-            task_id=task_id,
             run_id=run_id,
-            session=session,
+            task_id=task_id,
+            dag_id=dag_id,
             map_index=map_index,
         )
+        session.add(new)
+        session.flush()
     except TypeError as e:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -246,6 +268,34 @@ def set_xcom(
     return {"message": "XCom successfully set"}
 
 
+@router.delete(
+    "/{dag_id}/{run_id}/{task_id}/{key}",
+    responses={status.HTTP_404_NOT_FOUND: {"description": "XCom not found"}},
+    description="Delete a single XCom Value",
+)
+def delete_xcom(
+    session: SessionDep,
+    token: deps.TokenDep,
+    dag_id: str,
+    run_id: str,
+    task_id: str,
+    key: str,
+):
+    if not has_xcom_access(dag_id, run_id, task_id, key, token, write=True):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail={
+                "reason": "access_denied",
+                "message": f"Task does not have access to delete XCom with key '{key}'",
+            },
+        )
+
+    query = session.query(XComModel).where(XComModel.key == key).first()
+    session.delete(query)
+    session.commit()
+    return {"message": f"XCom with key: {key} successfully deleted."}
+
+
 def has_xcom_access(
     dag_id: str, run_id: str, task_id: str, xcom_key: str, token: TIToken, write: bool = False
 ) -> bool:

airflow/config_templates/config.yml

@@ -401,7 +401,7 @@ core:
       version_added: 1.10.12
       type: string
       example: "path.to.CustomXCom"
-      default: "airflow.models.xcom.BaseXCom"
+      default: "airflow.sdk.execution_time.xcom.BaseXCom"
     lazy_load_plugins:
       description: |
         By default Airflow plugins are lazily-loaded (only loaded when required). Set it to ``False``,