Docker-airflow: Airflow: task after BranchPythonOperator Task getting skipped

[manishgupta24]

I created a BranchPythonOperator which calls 2 tasks depending on the condition like:

typicon_check_table = BranchPythonOperator(
    task_id='typicon_check_table',
    python_callable=CheckTable(),
    provide_context=True,
    dag=typicon_task_dag)

typicon_create_table = PythonOperator(
    task_id='typicon_create_table',
    python_callable=CreateTable(),
    provide_context=True,
    dag=typicon_task_dag)

typicon_load_data = PythonOperator(
    task_id='typicon_load_data',
    python_callable=LoadData(),
    provide_context=True,
    dag=typicon_task_dag)

typicon_check_table.set_downstream([typicon_load_data, typicon_create_table])
typicon_create_table.set_downstream(typicon_load_data)

This is the CheckTable callable class:

class CheckTable:
    """
    DAG task to check if table exists or not.
    """

    def __call__(self, **kwargs) -> None:
        pg_hook = PostgresHook(postgres_conn_id="postgres_docker")
        query = "SELECT EXISTS ( \
            SELECT 1 FROM information_schema.tables \
            WHERE table_schema = 'public' \
            AND table_name = 'users');"

        table_exists = pg_hook.get_records(query)[0][0]
        if table_exists:
            return "typicon_load_data"
        return "typicon_create_table"

The issue is both the tasks are getting skipped when the typicon_check_table task is run.

How to fix this issue?

[zhongjiajie]

@manishgupta24 Mabe due to task dependence

typicon_check_table.set_downstream([typicon_load_data, typicon_create_table])
typicon_create_table.set_downstream(typicon_load_data)

If typicon_check_table return typicon_load_data then typicon_create_table would skip and downstream typicon_load_data would be skip too.

I you want to run typicon_load_data in any situation, you can do something like below

class CheckTable:
    """
    DAG task to check if table exists or not.
    """

    def __call__(self, **kwargs) -> None:
        pg_hook = PostgresHook(postgres_conn_id="postgres_docker")
        query = "SELECT EXISTS ( \
            SELECT 1 FROM information_schema.tables \
            WHERE table_schema = 'public' \
            AND table_name = 'users');"

        table_exists = pg_hook.get_records(query)[0][0]
        if table_exists:
>>> CHANGE           return "not_create_table"
        return "typicon_create_table"

typicon_check_table = BranchPythonOperator(
    task_id='typicon_check_table',
    python_callable=CheckTable(),
    provide_context=True,
    dag=typicon_task_dag)

typicon_create_table = PythonOperator(
    task_id='typicon_create_table',
    python_callable=CreateTable(),
    provide_context=True,
    dag=typicon_task_dag)

typicon_load_data = PythonOperator(
    task_id='typicon_load_data',
    python_callable=LoadData(),
    provide_context=True,
    dag=typicon_task_dag)

>>> CHANGE create two new task
not_create_table = DummyOperator(
    task_id='not_create_table',
    dag=typicon_task_dag
)

one_sucess = DummyOperator(
    task_id='one_task_sucess',
    trigger_rule=TriggerRule.ONE_SUCESS,
    dag=typicon_task_dag
)
>>> CHANGE

typicon_check_table.set_downstream([not_create_table, typicon_create_table])
one_sucess.set_upstream([not_create_table, typicon_create_table])
one_sucess.set_downstream(typicon_load_data)

[manishgupta24]

Thanks. Instead of using the dummy operators I used the TriggerRule.ONE_SUCCESS directly on typicon_load_data task because that has to run in any case every time the DAG is run. Works perfectly fine now.

[zhongjiajie]

@manishgupta24 that nice