An exmpale DAG testing the new SubDagOperator

example_rewrite_subdag.py

from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.subdag_operator import SubDagOperator
from airflow.utils.dates import days_ago


def dummy_dag(dag_name, args):
    dag_subdag = DAG(
        dag_id=dag_name,
        default_args=args,
        schedule_interval="@daily",
    )

    dummies = [DummyOperator(
            task_id='%s-task-%s' % (dag_name, i + 1),
            default_args=args,
            dag=dag_subdag,
        ) for i in range(3)]
    return dag_subdag


def subdag(dag_name, args):
    """
    Generate a DAG to be used as a subdag.

    :param str dag_name: Id of the DAG
    :param dict args: Default arguments to provide to the subdag
    :return: DAG to use as a subdag
    :rtype: airflow.models.DAG
    """
    dag_subdag = DAG(
        dag_id=dag_name,
        default_args=args,
        schedule_interval="@daily",
    )
    
    dummies = [DummyOperator(
            task_id='%s-task-%s' % (dag_name, i + 1),
            default_args=args,
            dag=dag_subdag,
        ) for i in range(5)]

    inside_section_1 = SubDagOperator(
        task_id=f'{dag_name}.inside-section-1',
        subdag_factory=dummy_dag,
        subdag_args=[args],
        dag=dag_subdag,
    )

    inside_section_2 = SubDagOperator(
        task_id=f'{dag_name}.inside-section-2',
        subdag_factory=dummy_dag,
        subdag_args=[args],
        dag=dag_subdag,
    )

    dummies[-1] >> inside_section_1
    dummies[-2] >> inside_section_2
    return dag_subdag

DAG_NAME = 'test_subdag'

args = {
    'owner': 'airflow',
    'start_date': days_ago(2),
}

dag = DAG(
    dag_id=DAG_NAME,
    default_args=args,
    schedule_interval="@once",
    tags=['example']
)

start = DummyOperator(
    task_id='start',
    dag=dag,
)

section_1 = SubDagOperator(
    task_id='section-1',
    subdag_factory=subdag,
    subdag_args=[args],
    dag=dag,
)

some_other_task = DummyOperator(
    task_id='some-other-task',
    dag=dag,
)

section_2 = SubDagOperator(
    task_id='section-2',
    subdag_factory=subdag,
    subdag_args=[args],
    dag=dag,
)

end = DummyOperator(
    task_id='end',
    dag=dag,
)

start >> section_1 >> some_other_task
some_other_task >> section_2 >> end