AIP-58: Add object storage backend for xcom (#37058)

* Add object storage backend for xcom

This adds the possibility to store xcoms on a  configurable object storage
supported backend. 

---------

Co-authored-by: Ephraim Anierobi <splendidzigy24@gmail.com>

airflow/models/xcom.py

@@ -29,6 +29,7 @@
 from typing import TYPE_CHECKING, Any, Generator, Iterable, cast, overload
 
 import attr
+from deprecated import deprecated
 from sqlalchemy import (
     Column,
     ForeignKeyConstraint,
@@ -368,6 +369,7 @@ def get_one(
     @staticmethod
     @provide_session
     @internal_api_call
+    @deprecated
     def get_one(
         execution_date: datetime.datetime | None = None,
         key: str | None = None,
@@ -418,7 +420,7 @@ def get_one(
 
         result = query.with_entities(BaseXCom.value).first()
         if result:
-            return BaseXCom.deserialize_value(result)
+            return XCom.deserialize_value(result)
         return None
 
     @overload
@@ -556,9 +558,15 @@ def delete(cls, xcoms: XCom | Iterable[XCom], session: Session) -> None:
         for xcom in xcoms:
             if not isinstance(xcom, XCom):
                 raise TypeError(f"Expected XCom; received {xcom.__class__.__name__}")
+            XCom.purge(xcom, session)
             session.delete(xcom)
         session.commit()
 
+    @staticmethod
+    def purge(xcom: XCom, session: Session) -> None:
+        """Purge an XCom entry from underlying storage implementations."""
+        pass
+
     @overload
     @staticmethod
     @internal_api_call
@@ -641,7 +649,13 @@ def clear(
         query = session.query(BaseXCom).filter_by(dag_id=dag_id, task_id=task_id, run_id=run_id)
         if map_index is not None:
             query = query.filter_by(map_index=map_index)
-        query.delete()
+
+        for xcom in query:
+            # print(f"Clearing XCOM {xcom} with value {xcom.value}")
+            XCom.purge(xcom, session)
+            session.delete(xcom)
+
+        session.commit()
 
     @staticmethod
     def serialize_value(

airflow/provider.yaml.schema.json

@@ -189,6 +189,13 @@
                 "type": "string"
             }
         },
+        "xcom": {
+            "type": "array",
+            "description": "XCom module names",
+            "items": {
+                "type": "string"
+            }
+        },
         "transfers": {
             "type": "array",
             "items": {

airflow/providers/common/io/provider.yaml

@@ -43,3 +43,37 @@ operators:
   - integration-name: Common IO
     python-modules:
       - airflow.providers.common.io.operators.file_transfer
+
+xcom:
+  - airflow.providers.common.io.xcom.backend
+
+config:
+  common.io:
+    description: Common IO configuration section
+    options:
+      xcom_objectstorage_path:
+        description: |
+          Path to a location on object storage where XComs can be stored in url format.
+        version_added: 1.3.0
+        type: string
+        example: "s3://conn_id@bucket/path"
+        default: ""
+      xcom_objectstorage_threshold:
+        description: |
+          Threshold in bytes for storing XComs in object storage. -1 means always store in the
+          database. 0 means always store in object storage. Any positive number means
+          it will be stored in object storage if the size of the value is greater than the threshold.
+        version_added: 1.3.0
+        type: integer
+        example: "1000000"
+        default: "-1"
+      xcom_objectstorage_compression:
+        description: |
+          Compression algorithm to use when storing XComs in object storage. Supported algorithms
+          are a.o.: snappy, zip, gzip, bz2, and lzma. If not specified, no compression will be used.
+          Note that the compression algorithm must be available in the Python installation (e.g.
+          python-snappy for snappy). Zip, gz, bz2 are available by default.
+        version_added: 1.3.0
+        type: string
+        example: "gz"
+        default: ""

airflow/providers/common/io/xcom/__init__.py

@@ -0,0 +1,32 @@
+# 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 __future__ import annotations
+
+import packaging.version
+
+try:
+    from airflow import __version__ as airflow_version
+except ImportError:
+    from airflow.version import version as airflow_version
+
+if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
+    "2.9.0"
+):
+    raise RuntimeError(
+        "The package xcom backend feature of `apache-airflow-providers-common-io` needs "
+        "Apache Airflow 2.9.0+"
+    )

airflow/providers/common/io/xcom/backend.py

@@ -0,0 +1,168 @@
+# 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 __future__ import annotations
+
+import json
+import uuid
+from typing import TYPE_CHECKING, Any, TypeVar
+from urllib.parse import urlsplit
+
+import fsspec.utils
+
+from airflow.configuration import conf
+from airflow.io.path import ObjectStoragePath
+from airflow.models.xcom import BaseXCom
+from airflow.utils.json import XComDecoder, XComEncoder
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+    from airflow.models import XCom
+
+T = TypeVar("T")
+
+SECTION = "common.io"
+
+
+def _is_relative_to(o: ObjectStoragePath, other: ObjectStoragePath) -> bool:
+    """This is a port of the pathlib.Path.is_relative_to method. It is not available in python 3.8."""
+    if hasattr(o, "is_relative_to"):
+        return o.is_relative_to(other)
+
+    try:
+        o.relative_to(other)
+        return True
+    except ValueError:
+        return False
+
+
+def _get_compression_suffix(compression: str) -> str:
+    """This returns the compression suffix for the given compression.
+
+    :raises ValueError: if the compression is not supported
+    """
+    for suffix, c in fsspec.utils.compressions.items():
+        if c == compression:
+            return suffix
+
+    raise ValueError(f"Compression {compression} is not supported. Make sure it is installed.")
+
+
+class XComObjectStoreBackend(BaseXCom):
+    """XCom backend that stores data in an object store or database depending on the size of the data.
+
+    If the value is larger than the configured threshold, it will be stored in an object store.
+    Otherwise, it will be stored in the database. If it is stored in an object store, the path
+    to the object in the store will be returned and saved in the database (by BaseXCom). Otherwise, the value
+    itself will be returned and thus saved in the database.
+    """
+
+    @staticmethod
+    def _get_key(data: str) -> str:
+        """This gets the key from the url and normalizes it to be relative to the configured path.
+
+        :raises ValueError: if the key is not relative to the configured path
+        :raises TypeError: if the url is not a valid url or cannot be split
+        """
+        path = conf.get(SECTION, "xcom_objectstore_path", fallback="")
+        p = ObjectStoragePath(path)
+
+        try:
+            url = urlsplit(data)
+        except AttributeError:
+            raise TypeError(f"Not a valid url: {data}")
+
+        if url.scheme:
+            k = ObjectStoragePath(data)
+
+            if _is_relative_to(k, p) is False:
+                raise ValueError(f"Invalid key: {data}")
+            else:
+                return data.replace(path, "", 1).lstrip("/")
+
+        raise ValueError(f"Not a valid url: {data}")
+
+    @staticmethod
+    def serialize_value(
+        value: T,
+        *,
+        key: str | None = None,
+        task_id: str | None = None,
+        dag_id: str | None = None,
+        run_id: str | None = None,
+        map_index: int | None = None,
+    ) -> bytes | str:
+        # we will always serialize ourselves and not by BaseXCom as the deserialize method
+        # from BaseXCom accepts only XCom objects and not the value directly
+        s_val = json.dumps(value, cls=XComEncoder).encode("utf-8")
+        path = conf.get(SECTION, "xcom_objectstore_path", fallback="")
+        compression = conf.get(SECTION, "xcom_objectstore_compression", fallback=None)
+
+        if compression:
+            suffix = "." + _get_compression_suffix(compression)
+        else:
+            suffix = ""
+
+        threshold = conf.getint(SECTION, "xcom_objectstore_threshold", fallback=-1)
+
+        if path and -1 < threshold < len(s_val):
+            # safeguard against collisions
+            while True:
+                p = ObjectStoragePath(path) / f"{dag_id}/{run_id}/{task_id}/{str(uuid.uuid4())}{suffix}"
+                if not p.exists():
+                    break
+
+            if not p.parent.exists():
+                p.parent.mkdir(parents=True, exist_ok=True)
+
+            with p.open("wb", compression=compression) as f:
+                f.write(s_val)
+
+            return BaseXCom.serialize_value(str(p))
+        else:
+            return s_val
+
+    @staticmethod
+    def deserialize_value(
+        result: XCom,
+    ) -> Any:
+        """Deserializes the value from the database or object storage.
+
+        Compression is inferred from the file extension.
+        """
+        data = BaseXCom.deserialize_value(result)
+        path = conf.get(SECTION, "xcom_objectstore_path", fallback="")
+
+        try:
+            p = ObjectStoragePath(path) / XComObjectStoreBackend._get_key(data)
+            return json.load(p.open("rb", compression="infer"), cls=XComDecoder)
+        except TypeError:
+            return data
+        except ValueError:
+            return data
+
+    @staticmethod
+    def purge(xcom: XCom, session: Session) -> None:
+        path = conf.get(SECTION, "xcom_objectstore_path", fallback="")
+        if isinstance(xcom.value, str):
+            try:
+                p = ObjectStoragePath(path) / XComObjectStoreBackend._get_key(xcom.value)
+                p.unlink(missing_ok=True)
+            except TypeError:
+                pass
+            except ValueError:
+                pass

docs/apache-airflow-providers-common-io/configurations-ref.rst

@@ -0,0 +1,18 @@
+ .. 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.
+
+.. include:: ../exts/includes/providers-configurations-ref.rst

docs/apache-airflow-providers-common-io/index.rst

@@ -28,6 +28,7 @@
     Home <self>
     Changelog <changelog>
     Security <security>
+    Configuration <configurations-ref>
 
 .. toctree::
     :hidden:
@@ -36,6 +37,7 @@
 
     Transferring a file <transfer>
     Operators <operators>
+    Object Storage XCom Backend <xcom_backend>
 
 .. toctree::
     :hidden:

docs/apache-airflow-providers-common-io/xcom_backend.rst

@@ -0,0 +1,42 @@
+.. 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.
+
+Object Storage XCom Backend
+===========================
+
+The default XCom backend is the :class:`~airflow.models.xcom.BaseXCom` class, which stores XComs in the Airflow database. This is fine for small values, but can be problematic for large values, or for large numbers of XComs.
+
+To enable storing XComs in an object store, you can set the ``xcom_backend`` configuration option to ``airflow.providers.common.io.xcom.backend.XComObjectStoreBackend``. You will also need to set ``xcom_objectstorage_path`` to the desired location. The connection
+id is obtained from the user part of the url the you will provide, e.g. ``xcom_objectstorage_path = s3://conn_id@mybucket/key``. Furthermore, ``xcom_objectstorage_threshold`` is required
+to be something larger than -1. Any object smaller than the threshold in bytes will be stored in the database and anything larger will be be
+put in object storage. This will allow a hybrid setup. If an xcom is stored on object storage a reference will be
+saved in the database. Finally, you can set ``xcom_objectstorage_compression`` to fsspec supported compression methods like ``zip`` or ``snappy`` to
+compress the data before storing it in object storage.
+
+So for example the following configuration will store anything above 1MB in S3 and will compress it using gzip::
+
+      [core]
+      xcom_backend = airflow.providers.common.io.xcom.backend.XComObjectStoreBackend
+
+      [common.io]
+      xcom_objectstorage_path = s3://conn_id@mybucket/key
+      xcom_objectstorage_threshold = 1048576
+      xcom_objectstorage_compression = gzip
+
+.. note::
+
+  Compression requires the support for it is installed in your python environment. For example, to use ``snappy`` compression, you need to install ``python-snappy``. Zip, gzip and bz2 work out of the box.