Repository: apache/airflow
Pull request: Allow to define custom XCom class
All details can be found in the following issue:
closes: #8059

---
Make sure to mark the boxes below before creating PR: [x]

- [x] Description above provides context of the change
- [x] Unit tests coverage for changes (not needed for documentation changes)
- [x] Target Github ISSUE in description if exists
- [x] Commits follow "[How to write a good git commit message](http://chris.beams.io/posts/git-commit/)"
- [x] Relevant documentation is updated including usage instructions.
- [x] I will engage committers as explained in [Contribution Workflow Example](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#contribution-workflow-example).

---
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
Read the [Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines) for more information.


Problem statement:
[AIP-31] Enable arbitrary XCom class

**Description**



- Enable custom XCom backend functionality for custom serialization/deserialization.
- XCom backend can be set from airflow.cfg and defaults to existing `airflow.models.XCom`. 



**Use case / motivation**
This comes from a discussion with @turbaszek . We were talking that it would be nice to have a flexible XCom backend that allowed us to serialize/deserialize XCom values using custom logic. A possible extension here would be to enable large objects to be serialized into a custom file system while XCom would only include a pointer (and metadata) of the object in the file system. 

An example use case would be data processing:

```

@task
def process_data(input: pd.DataFrame) -> pd.DataFrame
   # make some magic
```

Where DataFrame are saved between tasks as `csv`.

This is probably a stretch and maybe out of scope for this AIP (or require a new AIP). But wanted to put this out there.

Hints from issue discussion:
Thanks @casassg for putting this into an issue! 
 
The idea comes from DAGs where we have two operators OP1 >> OP2:
- OP1 downloads from API and upload file on GCS
- OP2 downloads the previous file from GCS and does something

Abstracting this layer of writing to and reading from persistent storage will make writing functional operators much easier. Because then my function will be called `process_data` not `download_and_process_data`.  That's why I think it's related to UX of functionals operators




+1 to Making Xcom backend pluggable with the default of `airflow.models.xcom.XCom`

Commit message: fixup! fixup! fixup! Allow to define custom XCom class

Patch to review:
diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 201af0c3e1150..4a4c248cf516d 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -359,6 +359,13 @@
       type: string
       example: ~
       default: "True"
+    - name: xcom_backend
+      description: |
+        Path to custom XCom class that will be used to store and resolve operators results
+      version_added: 2.0.0
+      type: string
+      example: "path.to.CustomXCom"
+      default: "airflow.models.xcom.BaseXCom"
 
 - name: logging
   description: ~
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index cc86e15cbefac..fe58671ce2e29 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -208,6 +208,10 @@ max_num_rendered_ti_fields_per_task = 30
 # On each dagrun check against defined SLAs
 check_slas = True
 
+# Path to custom XCom class that will be used to store and resolve operators results
+# Example: xcom_backend = path.to.CustomXCom
+xcom_backend = airflow.models.xcom.BaseXCom
+
 [logging]
 # The folder where airflow should store its log files
 # This path must be absolute
diff --git a/airflow/models/xcom.py b/airflow/models/xcom.py
index 86c4bfb6dea58..d68db89eab69f 100644
--- a/airflow/models/xcom.py
+++ b/airflow/models/xcom.py
@@ -41,7 +41,7 @@
 XCOM_RETURN_KEY = 'return_value'
 
 
-class XCom(Base, LoggingMixin):
+class BaseXCom(Base, LoggingMixin):
     """
     Base class for XCom objects.
     """
@@ -229,3 +229,16 @@ def deserialize_value(result) -> Any:
                       "for XCOM, then you need to enable pickle "
                       "support for XCOM in your airflow config.")
             raise
+
+
+def resolve_xcom_backend():
+    """Resolves custom XCom class"""
+    clazz = conf.getimport("core", "xcom_backend")
+    if not issubclass(clazz, BaseXCom):
+        raise TypeError(
+            f"Your custom XCom class `{clazz.__name__}` is not a subclass of `{BaseXCom.__name__}`."
+        )
+    return clazz
+
+
+XCom = resolve_xcom_backend()
diff --git a/docs/concepts.rst b/docs/concepts.rst
index 20eb7d41dc0e0..92e6e750f65ae 100644
--- a/docs/concepts.rst
+++ b/docs/concepts.rst
@@ -649,6 +649,15 @@ of what this may look like:
 Note that XComs are similar to `Variables`_, but are specifically designed
 for inter-task communication rather than global settings.
 
+Custom XCom backend
+'''''''''''''''''''
+
+It is possible to change ``XCom`` behaviour os serialization and deserialization of tasks' result.
+To do this one have to change ``xcom_backend`` parameter in Airflow config. Provided value should point
+to a class that is subclass of :class:`~airflow.models.xcom.BaseXCom`. To alter the serialaization /
+deserialization mechanism the custom class should override ``serialize_value`` and ``deserialize_value``
+methods.
+
 .. _concepts:variables:
 
 Variables
diff --git a/tests/models/test_xcom.py b/tests/models/test_xcom.py
new file mode 100644
index 0000000000000..590596bd85f1f
--- /dev/null
+++ b/tests/models/test_xcom.py
@@ -0,0 +1,36 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow.models import XCom
+from airflow.models.xcom import resolve_xcom_backend
+from tests.test_utils.config import conf_vars
+
+
+class CustomXCom(XCom):
+    @staticmethod
+    def serialize_value(_):
+        return "custom_value"
+
+
+class TestXCom:
+    def test_resolve_xcom_class(self):
+        with conf_vars(
+            {("core", "xcom_backend"): "tests.models.test_xcom.CustomXCom"}
+        ):
+            cls = resolve_xcom_backend()
+            assert issubclass(cls, CustomXCom)
+            assert cls().serialize_value(None) == "custom_value"


Code review comments:
diff_hunk: @@ -229,3 +229,16 @@ def deserialize_value(result) -> Any:
                       "for XCOM, then you need to enable pickle "
                       "support for XCOM in your airflow config.")
             raise
+
+
+def resolve_xcom_backend():
+    """Resolves custom XCom class"""
+    clazz = conf.getimport("core", "xcom_backend")
path: airflow/models/xcom.py
text: @user1:
should we default to `BaseXCom` here? Just in case, someone updates the code, but not the config.

@author:
Thanks @user1 , good catch! I fixed it and extended tests.