PythonVirtualenvOperator with provide_context=True does not have 'ti' keyword

Apache Airflow version: 1.10.13

Environment: Docker (Ubuntu 18.4 - Python 3.7)

• Cloud provider or hardware configuration: Azure
• OS (e.g. from /etc/os-release): Docker (Ubuntu 18.4 - Python 3.7)
• Kernel (e.g. uname -a): Docker (Ubuntu 18.4 - Python 3.7)
• Install tools: N/A
• Others: N/A

What happened:

When we enable provide_context=True for PythonVirtualenvOperator and try to use the xcom_push to pass a variable I get this error:

File "/tmp/venv0upgqome/script.py", line 13, in push\n kwargs[\'ti\'].xcom_push(key=\'value from pusher 1\', value=value_1)\nKeyError: \'ti\'\n'

How to reproduce it:

import airflow
from airflow import DAG
from airflow.operators.python_operator import PythonVirtualenvOperator

args = {
    'owner': 'Airflow',
    'start_date': airflow.utils.dates.days_ago(2),

}

dag = DAG('BAtatas', schedule_interval="@once", default_args=args)


def push(**kwargs):
    """Pushes an XCom without a specific target"""
    value_1 = [1, 2, 3]
    print("printing the kwargs!!!")
    print(kwargs)
    kwargs['ti'].xcom_push(key='value from pusher 1', value=value_1)


def push_by_returning(**kwargs):
    value_2 = {'a': 'b'}
    """Pushes an XCom without a specific target, just by returning it"""
    return value_2


def puller(**kwargs):
    value_2 = {'a': 'b'}
    value_1 = [1, 2, 3]
    """Pull all previously pushed XComs and check if the pushed values match the pulled values."""
    ti = kwargs['ti']

    # get value_1
    pulled_value_1 = ti.xcom_pull(key=None, task_ids='push')
    assert pulled_value_1 == value_1

    # get value_2
    pulled_value_2 = ti.xcom_pull(task_ids='push_by_returning')
    assert pulled_value_2 == value_2

    # get both value_1 and value_2
    pulled_value_1, pulled_value_2 = ti.xcom_pull(
        key=None, task_ids=['push', 'push_by_returning'])
    assert (pulled_value_1, pulled_value_2) == (value_1, value_2)


push1 = PythonVirtualenvOperator(
    task_id='push',
    dag=dag,
    python_callable=push,
    requirements=[],
    python_version='3.7',
    use_dill=False,
    provide_context=True,
    system_site_packages=True,

)

push2 = PythonVirtualenvOperator(
    task_id='push_by_returning',
    dag=dag,
    python_callable=push_by_returning,
    requirements=[],
    python_version='3.7',
    use_dill=False,
    provide_context=True,
    system_site_packages=True,

)

pull = PythonVirtualenvOperator(
    task_id='puller',
    dag=dag,
    python_callable=puller,
    requirements=[],
    python_version='3.7',
    use_dill=False,
    provide_context=True,
    system_site_packages=True,

)

pull << [push1, push2]

[uranusjr]

Feel free to propose a PR to add ti to the list.

[xinbinhuang]

Also found this Jira story (https://issues.apache.org/jira/browse/AIRFLOW-2738) that’s still open that seems to be addressing the same problem.

Just a quick note: the community has migrated the issues tracking system from Jira to GitHub issues, so whatever left in Jira right now is likely outdated or inaccurate.