Unable to pass xcom in Custom Operators in Airflow

user_11286659 0
Viewed 740

I have a simple, linear DAG(created using Airflow 2.0) with two tasks. I have custom operators for each of the task which extend over BaseOperator. Following is the code for dag and operators:-

class Operator1(BaseOperator):
    @apply_defaults
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
    def execute(self, context):
        ...
        logging.info('First task')
        context['task_instance'].xcom_push(key="payload", value=data)
        return data

class Operator2(BaseOperator):
    @apply_defaults
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
    def execute(self, context):
        ...
        logging.info("context is ", context)
        parameters = context['task_instance'].xcom_pull(key="payload", value=data)
 

with DAG('dag_1', default_args=DEFAULT_ARGS, schedule_interval=None) as dag:
    TASK_1 = Operator1(
        task_id='task_1',
        do_xcom_push=True)
    TASK_2 = Operator2(
        task_id='task_2',
        do_xcom_push=True)
    TASK_1 >> TASK_2

When I run the DAG, I find that context which is used for getting xcom values is empty. I have searched a lot of answers on stackoverflow and tried the way mentioned in them but they didn't work.

Would really appreciate some hint over the issue - how to push and pull xcom values in custom operators?

airflow
1 Answers

I took your code and run it, the first problem was that start_date wasn't defined, so it ended up in an exception:

Exception has occurred: AirflowException (note: full exception trace is shown but execution is paused at: _run_module_as_main)
Task is missing the start_date parameter

Also, in Operator1 class, data variable is not defined. I guess maybe you missed them when you made the code example.

Other than that the code worked, but I think you should consider defining the task_id parameter when doing the xcom_pull operation.

From TaskInstance xcom_pull method description:

:param task_ids: Only XComs from tasks with matching ids will be pulled. Can pass None to remove the filter.

Here is the code of a working example, note that I use two equivalent methods to perform the XComs operations:

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.utils.decorators import apply_defaults
from airflow.models import BaseOperator


class Operator1(BaseOperator):
    @apply_defaults
    def __init__(self,  *args, **kwargs) -> None:
        super(Operator1, self).__init__(*args, **kwargs)

    def execute(self, context):
        print('First task')
        data = "valuable_data"
        more_data = "more_valueable_data"
        context['task_instance'].xcom_push(key="payload", value=data)
        self.xcom_push(context, "more_data", more_data)
        return data


class Operator2(BaseOperator):
    @apply_defaults
    def __init__(self,  *args, **kwargs) -> None:
        super(Operator2, self).__init__(*args, **kwargs)

    def execute(self, context):
        # print(f"context is  {context}")
        data = context['task_instance'].xcom_pull(
            "task_1",
            key="payload")
        more_data = self.xcom_pull(context, "task_1", key="more_data")

        print(f"Obtained data: {data}")
        print(f"Obtained more_data: {more_data}")


with DAG('dag_1',
         default_args={'owner': 'airflow'},
         start_date=days_ago(1),
         catchup=False,
         schedule_interval=None) as dag:

    TASK_1 = Operator1(
        task_id='task_1'
    )
    TASK_2 = Operator2(
        task_id='task_2'
    )
    TASK_1 >> TASK_2

Log from Task_2:

[2021-06-15 12:55:01,206] {taskinstance.py:1255} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=dag_1
AIRFLOW_CTX_TASK_ID=task_2
AIRFLOW_CTX_EXECUTION_DATE=2021-06-14T00:00:00+00:00
AIRFLOW_CTX_DAG_RUN_ID=backfill__2021-06-14T00:00:00+00:00
Obtained data: valuable_data
Obtained more_data: more_valueable_data
[2021-06-15 12:55:01,227] {taskinstance.py:1159} INFO - Marking task as SUCCESS. dag_id=dag_1, task_id=task_2, execution_date=20210614T000000, start_date=20210615T120402, end_date=20210615T125501

Side notes: I changed the __init__ method in order to accept *args as well. I'm using print but It could be done using Airflow logger as self.log.info('msg') .

Let me know if that worked for you!

user_105692200
Related