Add CRUD endpoint for XCom

airflow/api_connexion/endpoints/xcom_endpoint.py

@@ -14,25 +14,47 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from flask import request
+from typing import List, Optional
+
+from flask import Response, request
+from marshmallow import ValidationError
 from sqlalchemy import and_, func
 from sqlalchemy.orm.session import Session
 
 from airflow.api_connexion import parameters
-from airflow.api_connexion.exceptions import NotFound
+from airflow.api_connexion.exceptions import BadRequest, NotFound
 from airflow.api_connexion.schemas.xcom_schema import (
     XComCollection, XComCollectionItemSchema, XComCollectionSchema, xcom_collection_item_schema,
-    xcom_collection_schema,
+    xcom_collection_schema, xcom_schema,
 )
 from airflow.models import DagRun as DR, XCom
 from airflow.utils.session import provide_session
 
 
-def delete_xcom_entry():
+@provide_session
+def delete_xcom_entry(
+    dag_id: str,
+    dag_run_id: str,
+    task_id: str,
+    xcom_key: str,
+    session: Session
+):
     """
     Delete an XCom entry
     """
-    raise NotImplementedError("Not implemented yet.")
+
+    dag_run = session.query(DR).filter(DR.run_id == dag_run_id,
+                                       DR.dag_id == dag_id).first()
+    if not dag_run:
+        raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found')
+
+    entry = session.query(XCom).filter(XCom.dag_id == dag_id,
+                                       XCom.task_id == task_id,
+                                       XCom.key == xcom_key).delete()
+    if not entry:
+        raise NotFound("XCom entry not found")
+
+    return Response('', 204)
 
 
 @provide_session
@@ -89,15 +111,78 @@ def get_xcom_entry(
     return xcom_collection_item_schema.dump(query_object)
 
 
-def patch_xcom_entry():
+@provide_session
+def patch_xcom_entry(
+    session: Session,
+    dag_id: str,
+    dag_run_id: str,
+    task_id: str,
+    xcom_key: str,
+    update_mask: Optional[List[str]] = None
+):
     """
     Update an XCom entry
     """
-    raise NotImplementedError("Not implemented yet.")
+    dag_run = session.query(DR).filter(DR.run_id == dag_run_id,
+                                       DR.dag_id == dag_id).first()
+    if not dag_run:
+        raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found')
+
+    xcom_entry = session.query(XCom).filter(XCom.dag_id == dag_id,
+                                            XCom.task_id == task_id,
+                                            XCom.key == xcom_key).first()
+
+    if not xcom_entry:
+        raise NotFound("XCom not found")
+    try:
+        body = xcom_schema.load(request.json, partial=("dag_id", "task_id",
+                                                       "key", "execution_date"))
+    except ValidationError as err:
+        raise BadRequest(detail=err.messages.get('_schema', [str(err.messages)])[0])
+    data = body.data
+    # Check that other attributes are not being updated
+    for field in data:
+        if field != 'value' and data[field] != getattr(xcom_entry, field):
+            raise BadRequest(detail=f"'{field}' cannot be updated")
+    if update_mask:
+        update_mask = [i.strip() for i in update_mask]
+        data_ = {}
+        for field in update_mask:
+            if field in data and field != 'value':
+                raise BadRequest(detail=f"{field} cannot be updated")
+            elif field in data and field == 'value':
+                data_[field] = data[field]
+            else:
+                raise BadRequest(detail=f"Unknown field '{field}' in update_mask")
+        data = data_
+    data['value'] = XCom.serialize_value(data['value'])
+    xcom_entry.value = data['value']
+    session.add(xcom_entry)
+    session.commit()
+    return xcom_schema.dump(xcom_entry)
 
 
-def post_xcom_entries():
+@provide_session
+def post_xcom_entries(dag_id: str,
+                      dag_run_id: str,
+                      task_id: str,
+                      session: Session
+                      ):
     """
     Create an XCom entry
     """
-    raise NotImplementedError("Not implemented yet.")
+
+    dag_run = session.query(DR).filter(DR.run_id == dag_run_id, DR.dag_id == dag_id).first()
+    if not dag_run:
+        raise NotFound(f'DAGRun with dag_id:{dag_id} and run_id:{dag_run_id} not found')
+    try:
+        xcom_entry = xcom_schema.load(request.json)
+    except ValidationError as err:
+        raise BadRequest(detail=err.messages.get('_schema', [str(err.messages)])[0])
+    session.query(XCom).filter(XCom.dag_id == dag_id, XCom.task_id == task_id).delete()
+    data = xcom_entry.data
+    data['value'] = XCom.serialize_value(data['value'])
+    xcom = XCom(**data)
+    session.add(xcom)
+    session.commit()
+    return xcom_schema.dump(xcom)

airflow/api_connexion/schemas/xcom_schema.py

@@ -16,7 +16,7 @@
 # under the License.
 from typing import List, NamedTuple
 
-from marshmallow import Schema, fields
+from marshmallow import Schema, ValidationError, fields, validates_schema
 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
 
 from airflow.models import XCom
@@ -37,13 +37,20 @@ class Meta:
     task_id = auto_field()
     dag_id = auto_field()
 
+    @validates_schema(pass_original=True)
+    def check_unknown_fields(self, data, original_data):
+        """ Validates unknown field """
+        unknown = set(original_data) - set(self.fields)
+        if unknown:
+            raise ValidationError(f'Extra arguments passed: {list(unknown)}')
+
 
 class XComSchema(XComCollectionItemSchema):
     """
     XCom schema
     """
 
-    value = auto_field()
+    value = auto_field(required=True)
 
 
 class XComCollection(NamedTuple):