Using Airflow xComm custom backend to share data

Instructions

Let’s dive into an actual setup of the backend. 

1. Create an include directory in your Airflow default directory. In that directory you will need to create a file for the custom xComm backend code. In our case it is called s3_xcom_backend.py

Inside that file you can post the ready made code listed bellow.

import uuid
import pandas as pd

from typing import Any
from airflow.models.xcom import BaseXCom
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
import pickle

class S3XComBackend(BaseXCom):
    PREFIX = "xcom_s3://"
    BUCKET_NAME = "{BUCKET_NAME}"

    @staticmethod
    def serialize_value(value: Any):
        print(f"DEBUG: Custom xComm beckend object of type {type(value)} with content: {value}")
        if isinstance(value, pd.DataFrame):
            print(f"DEBUG: serializing value of type {type(value)} with content: {value}")
            hook        = S3Hook()
            key         = "data_csv_" + str(uuid.uuid4())
            filename    = f"{key}.csv"

            value.to_csv(filename)
            hook.load_file(
                filename=filename,
                key=key,
                bucket_name=S3XComBackend.BUCKET_NAME,
                replace=True
            )
            value = S3XComBackend.PREFIX + key
            return BaseXCom.serialize_value(value)

        if isinstance(value, list) or isinstance(value, dict):
            print(f"DEBUG: serializing value of type {type(value)} with content: {value}")
            hook        = S3Hook()
            key         = "data_dat_" + str(uuid.uuid4())
            filename    = f"{key}.dat"

            with open(filename, 'wb') as f:
                pickle.dump(value, f)

            hook.load_file(
                filename=filename,
                key=key,
                bucket_name=S3XComBackend.BUCKET_NAME,
                replace=True
            )
            value = S3XComBackend.PREFIX + key
            return BaseXCom.serialize_value(value)

    @staticmethod
    def deserialize_value(result) -> Any:
        result = BaseXCom.deserialize_value(result)
        print("BaseXCOM deserialized result:", result)
        if isinstance(result, str) and result.startswith(S3XComBackend.PREFIX):
            hook    = S3Hook()
            key     = result.replace(S3XComBackend.PREFIX, "")
            filename = hook.download_file(
                key=key,
                bucket_name=S3XComBackend.BUCKET_NAME,
                preserve_file_name=True,
                local_path="/tmp"
            )

            print("filename:", filename)
            # get the file format to distinguish between different types of data
            fmt = filename.split('/')[3].split('_')[1]

            print("fmt:", fmt)

            if fmt == "csv":
                result = pd.read_csv(filename)
                return result

            if fmt == "dat":
                with open(filename, 'rb') as f:
                    result = pickle.load(f)
                return result