combine BranchPythonOperator and PythonVirtualenvOperator

Question

I have a PythonVirtualenvOperator which reads some data from a database - if there is no new data, then the DAG should end there, otherwise it should call additional tasks e.g

#dag.py

load_data >>[if_data,if_no_data]>>another_task>>last_task

I understand that it can be done using PythonBranchOperator but I can't see how I can combine the venv and the branch-operator.

Is it doable?

The PythonVirtualenvOperator is the load_data ? have you considered returning true/false with xcom ? — Elad Kalif
– Elad Kalif, Commented May 20, 2021 at 13:07

Elad Kalif · Accepted Answer · 2021-05-20 13:15:34Z

1

This can be solved using Xcom. load_date can push the number of records it processed (new data).

Your pipe can be:

def choose(**context):
    value = context['ti'].xcom_pull(task_ids='load_data')
    if int(value)>0:
        return 'if_data'
    return 'if_no_data'

branch = BranchPythonOperator(
    task_id='branch_task',
    provide_context=True, # Remove this line if Airflow>=2.0.0
    python_callable=choose)

load_data >> branch >>[if_data,if_no_data]>>another_task>>last_task

answered May 20, 2021 at 13:15

Elad Kalif

16.4k3 gold badges29 silver badges63 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

CutePoison Over a year ago

What if I want to stop when if_no_data is triggered i.e another_task and last_task should only be run after if_data

CutePoison Over a year ago

Nvm, I figured out adding another_task.set_upstream(if_data) and last_task.set_upstream(another_task) before load_data >> branch >>[if_data, no_data] made the deal

Abouklila · Accepted Answer · 2023-11-22 18:40:25Z

0

This can be solved by using BranchPythonVirtualenvOperator

    branching_venv = BranchPythonVirtualenvOperator(
        task_id="branching_venv",
        requirements=["numpy~=1.24.4"],
        python_callable=branch_with_venv
    )

answered Nov 22, 2023 at 18:40

Abouklila

133 bronze badges

Collectives™ on Stack Overflow

combine BranchPythonOperator and PythonVirtualenvOperator

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related