How to not display a task in DAG based on a condition - airflow

Question

I have a DAG and have 3 tasks in it. I would like to not display the 2nd task (middle_name) in the DAG run based on a condition. for e.g. if middle_name_var == 'false', i don't want to display the middle_name task in the DAG. Is there a way to elegantly achieve this?
from airflow.operators import PythonOperator
from airflow.operators.python_operator import BranchPythonOperator
from airflow.operators.dummy_operator import DummyOperator
from datetime import datetime, timedelta
from airflow.models import Variable
middle_name_var = Variable.get('middle_name')
default_args = {
'owner': 'test',
'depends_on_past': False,
'start_date': datetime(2018, 6, 18),
'email': ['tes#abc.com'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=1)
}
dag = DAG(
'name',
default_args=default_args,
schedule_interval="#once")
def first_name():
print('John')
def middle_name():
print('Smith')
def last_name():
print('Doe')
first_name_task = PythonOperator(
task_id='first_name',
provide_context=False,
python_callable=first_name,
dag=dag
)
middle_name_task = PythonOperator(
task_id='middle_name',
provide_context=False,
python_callable=middle_name,
dag=dag
)
last_name_task = PythonOperator(
task_id='last_name',
provide_context=False,
python_callable=last_name,
dag=dag
)
if middle_name_var == 'true':
first_name_task >> middle_name_task >>last_name_task
else:
first_name_task >> last_name_task
MY DAG looks like this with middle_name task...But i would like to not have middle_name task based on middle_name_var, which is set to false in this case.

Answer 1

With reference to that last set of chaining statements
# by the way i believe the comparison expression should be
# middle_name_var == True (boolean rather than string), but lets ignore it for now
if middle_name_var == 'true':
first_name_task >> middle_name_task >>last_name_task
else:
first_name_task >> last_name_task
Let me ask you: what would happen if you remove these chaining statements? Would the tasks disappear from DAG?
Not quite.
Chaining merely establishes a dependency relationship between tasks. Even without chaining, your task would remain part of your DAG (like in the screenshot that you've posted).
Here's the secret bit: a task becomes part of your dag as soon as you declare it
middle_name_task = PythonOperator(
task_id='middle_name',
provide_context=False,
python_callable=middle_name,
dag=dag
)
And whether or not you set that task upstream or downstream of some other tasks, it will continue to 'appear' in your DAG. Quoting the docs in this regard
Operators do not have to be assigned to DAGs immediately (previously
dag was a required argument). However, once an operator is assigned to
a DAG, it can not be transferred or unassigned. DAG assignment can be
done explicitly when the operator is created, through deferred
assignment, or even inferred from other operators.
Q So what should you do to 'not display' the task?
A Simply not declare (instantiate) it.
Q And how would you go about doing that?
A Just move task declaration inside your if-else clause
if middle_name_var == 'true':
middle_name_task = PythonOperator(
task_id='middle_name',
provide_context=False,
python_callable=middle_name,
dag=dag
)
first_name_task >> middle_name_task >>last_name_task
else:
first_name_task >> last_name_task