1
votes

1. Background

I created two custom operators. PushOperator is responsible for using xcom_push to write random numbers to key=batch_id, and PullOperator is responsible for using xcom_pull to retrieve the value of key=batch_id from the upstream PushOperator task.

my_operator.py

import random

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults


class PushOperator(BaseOperator):

    @apply_defaults
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.batch_id = random.randint(1, 100)

    def execute(self, context):
        self.xcom_push(context=context, key='batch_id', value=self.batch_id)


class PullOperator(BaseOperator):

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

    def execute(self, context):
        if self.upstream_task_ids is not None:
            for task_id in self.upstream_task_ids:
                batch_id = self.xcom_pull(context=context, task_ids=task_id, key='batch_id')
                self.log.info("===> PullOperator::TaskId={}, BatchId={}".format(task_id, batch_id))

my_dag.py

from datetime import datetime, timedelta
from airflow import DAG

from custom_operator.my_operator import PullOperator, PushOperator

# define DAG
default_args = {
    'owner': 'yewei.oyyw',
    'start_date': datetime(2020, 12, 10),
    'retries': 1,
    'retry_delay': timedelta(minutes=1),
}
dag = DAG(dag_id='my_dag', default_args=default_args, schedule_interval=timedelta(minutes=10))

# Task 1
Task1 = PushOperator(
    task_id="Task1",
    dag=dag)

# Task 2
Task2 = PushOperator(
    task_id="Task2",
    dag=dag)

# Task 3
Task3 = PushOperator(
    task_id="Task3",
    dag=dag)

# Task 4
PullTask = PullOperator(
    task_id="PullTask",
    dag=dag)

[Task1, Task2, Task3] >> PullTask

2. Question

Will the value of xcom be isolated based on Job, such as scheduling the same DAG multiple times, will the batch_id value written by PushOperator Job N overwrite the batch_id value written by Job 1

for example:
[Job1] Task1 -> batch_id = 1
[Job1] Task2 -> batch_id = 2
[Job1] Task3 -> batch_id = 3

[Job2] Task1 -> batch_id = 4
[Job2] Task2 -> batch_id = 5
[Job2] Task3 -> batch_id = 6

expect:
Job1 ===> PullOperator::TaskId=Task1, BatchId=1
Job1 ===> PullOperator::TaskId=Task2, BatchId=2
Job1 ===> PullOperator::TaskId=Task3, BatchId=3

Don't expect:
Job1 ===> PullOperator::TaskId=Task1, BatchId=4 (overridden by Job2 Task1)
Job1 ===> PullOperator::TaskId=Task2, BatchId=5 (overridden by Job2 Task2)
Job1 ===> PullOperator::TaskId=Task3, BatchId=6 (overridden by Job2 Task3)
1

1 Answers

1
votes

Yes It will be isolated XCom is based on a per job basis