xcom_pull not working in a scheduled DAG when triggered manually

[falcon97]

Apache Airflow version

Other Airflow 2 version (please specify below)

If "Other Airflow 2 version" selected, which one?

2.7.3

What happened?

I have a primary DAG which generates some timestamps during each DAG run. Following the completion of some tasks in this DAG, I trigger a few other DAGs using the TriggerDagRunOperator. These secondary DAGs make use of the timestamps generated by the primary one using the xcom_pull function. There was one instance where I had to trigger the DAG manually outside of its schedule. All the tasks completed as expected during this manual run, and fetched the values from XCom correctly.

But during the next scheduled run, I observed that in the secondary DAGs, the value fetched by the xcom_pull function was from the previous manual run of the primary DAG. However in the primary DAG itself, the value was being fetched correctly.

What you think should happen instead?

The xcom_pull function should be fetching the latest value regardless of whether the previous run is manual or scheduled. This is exhibiting the issue mentioned in #28085.
Relevant comment link: #28085 (comment)

How to reproduce

1. Create the DAGs mentioned below.
2. Unpause both DAGs. The first scheduled run of main_dag should push a value to XCom, and the corresponding DAG run of child_dag should fetch the same value from XCom.
3. Before the next scheduled run of main_dag begins in 5 minutes, trigger it manually. Again, we should see the same behavior as mentioned in step 2.
4. Let the next scheduled run of main_dag start. It pushes a value to XCom, but the same value is not fetched by child_dag. Instead, the value from the previous manual run is printed.

main_dag.py:

from datetime import datetime, timedelta
import os
import pendulum
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.trigger_dagrun import TriggerDagRunOperator 

airflow_include_dir = os.path.join(os.environ['AIRFLOW_HOME'], 'include')

default_args = {
    "owner": "airflow",
    'retries': 3,
    'retry_delay': timedelta(seconds=60),
}
 
def push_config(ti):
    current_datetime = datetime.now()
    ti.xcom_push(key='snapshot_id', value=current_datetime)
    print(f'Pushed {current_datetime} to XCom with key snapshot_id.')

 with DAG('main_dag', start_date=datetime(2023, 11, 8, tzinfo=pendulum.timezone("America/New_York")),
         schedule='*/5 * * * *', max_active_runs=1, catchup=False, template_searchpath=airflow_include_dir,
         default_args=default_args, tags=["basic"]) as dag:

    push_config = PythonOperator(
        task_id='push_config',
        python_callable=push_config
    )

    child_dag_trigger = TriggerDagRunOperator(
        task_id="child_dag_trigger",
        trigger_dag_id="child_dag",
        retries=0,
        wait_for_completion=True,
        deferrable=True
    )

    push_config >> child_dag_trigger

child_dag.py:

from airflow import DAG
from datetime import datetime, timedelta
from airflow.operators.python import PythonOperator
 
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2023, 11, 8),
    'retries': 3,
    'retry_delay': timedelta(seconds=60)
}

def pull_config(ti):
    value = ti.xcom_pull(
        dag_id='main_dag', task_ids='push_config', key='snapshot_id', include_prior_dates=True)
    print(value)

with DAG('child_dag', default_args=default_args, schedule=None, catchup=False, tags=["reporting"]) as dag:

    pull_config = PythonOperator(
        task_id='pull_config',
        python_callable=pull_config
    )

    pull_config

Operating System

Linux

Versions of Apache Airflow Providers

No response

Deployment

Astronomer

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[karenbraganz]

This is happening because the XCom is pulled from the DAG run with the most recent logical date when include_prior_dates=True. Here are two potential workarounds for this issue:

1. Pass the parent DAG's run ID to the child DAG through TriggerDagRunOperator and set include_prior_dates=False.

First, add the trigger_run_id="{{ run_id }}" parameter to TriggerDagRunOperator in main_dag.py as shown below.

    child_dag_trigger = TriggerDagRunOperator(
        task_id="child_dag_trigger",
        trigger_dag_id="child_dag",
        retries=0,
        wait_for_completion=True,
        deferrable=True
        trigger_run="{{ run_id }}"

Next, delete the include_prior_dates=True parameter from the XCom pull in child_dag.py as shown below. It will default to False.

value = ti.xcom_pull(dag_id='main_dag', task_ids='push_config', key='snapshot_id',)

When include_prior_dates=False, the child DAG run only pulls XComs from a DAG run with the same run_id. Since we are passing the parent DAG run's run_id to the child DAG run through TriggerDagRunOperator, the correct XCom will be pulled even during a DAG run that is scheduled after a manual run.

2. The second potential workaround would be to use CronTriggerTimetable in the DAG's schedule parameter instead of just the regular Cron expression which defaults to using the CronDataIntervalTimetable.

First, you would have to import CronTriggerTimetable by adding this line to main_dag.py.

from airflow.timetables.trigger import CronTriggerTimetable

Next, you would have to change the schedule parameter to schedule=CronTriggerTimetable(<your_Cron_expression>) as shown below in main_dag.py.

 with DAG('main_dag', start_date=datetime(2023, 11, 8, tzinfo=pendulum.timezone("America/New_York")),
         schedule=CronTriggerTimetable('*/5 * * * *', timezone='UTC'), max_active_runs=1, catchup=False, template_searchpath=airflow_include_dir,
         default_args=default_args, tags=["basic"]) as dag:

This works with include_prior_dates=True.

With CronTriggerTimetable, both manually triggered and scheduled DAG runs will have the logical_date equal to data_interval_end (time when DAG actually runs). On the other hand, with the default CronDataIntervalTimetable that we were using before, the logical_date equals the data_interval_end for manual runs only whereas the logical_date equals the data_interval_start for scheduled runs. This results in the logical_date of a scheduled run being prior to the logical_date of the manual run that was triggered right before it. Since the XCom pulls data from the run with the most recent logical_date when include_prior_dates=True, the child DAG run will pull the XCom from the manual run instead of its own parent run (which has a logical_date less recent than the manual run). Switching to CronTriggerTimetable fixes the logical_date logic and resolves this issue even when include_prior_dates=True.

[falcon97]

Thank you @karenbraganz for the inputs!
I followed the second workaround with CronTriggerTimetable and this resolves the problem.