Extending SubDAGOperator doesn't show Zoom into SubDAG in the UI

Apache Airflow version

2.1.x

Operating System

Debian GNU/Linux 10 (buster)

Versions of Apache Airflow Providers

No response

Deployment

Official Apache Airflow Helm Chart

Deployment details

Breeze

What happened

I extended the SubDAGOperator to add conf to template_fields as a workaround for #18491.

Extending the Zoom into SubDAG option in the UI doesn’t appear.

What you expected to happen

Zoom into SubDAG option should be available for Extended SubDAGOperators

How to reproduce

1. Use the following DAG

Click to expand!

# [START example_subdag_operator]
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.subdag import SubDagOperator
from airflow.utils.dates import days_ago
from datetime import datetime

DAG_NAME = 'example_subdag_operator'

args = {
    'owner': 'airflow',
}


class CustomSubDagOperator(SubDagOperator):
    template_fields = SubDagOperator.template_fields + ("conf",)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._task_type = 'SubDagOperator'

    @property
    def task_type(self):
        return self._task_type

def subddag_python(dag_run, **kwargs):
    print (kwargs)
    print (f"{dag_run.dag_id} {dag_run.conf}")

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

    :param str parent_dag_name: Id of the parent DAG
    :param str child_dag_name: Id of the child 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=f'{parent_dag_name}.{child_dag_name}',
        default_args=args,
        start_date=datetime(2021, 1, 1),
        catchup=False,
        schedule_interval="@daily",
    )

    tn = PythonOperator(
            task_id=f'subdag_python',
            python_callable=subddag_python,  # make sure you don't include the () of the function
            provide_context=True,
            dag=dag_subdag
    )

    return dag_subdag



def simple_python(dag_run):
    print (f"{dag_run.dag_id} {dag_run.conf}")

with DAG(
    dag_id=DAG_NAME, default_args=args, start_date=days_ago(2), schedule_interval="@once", tags=['example']
) as dag:

    tn = PythonOperator(
            task_id=f'simple_python',
            python_callable=simple_python,  # make sure you don't include the () of the function
            provide_context=True,
    )

    section_1 = CustomSubDagOperator(
        task_id='section-1',
        subdag=subdag(DAG_NAME, 'section-1', args),
        # conf="{{ dag_run.conf }}"
    )

    section_2 = SubDagOperator(
        task_id='section-2',
        subdag=subdag(DAG_NAME, 'section-2', args),
        # conf="{{ dag_run.conf }}"
    )

    tn >> section_1
    tn >> section_2

1. On the airflow UI you will see Zoom into SubDAG option for the direct task but not for the extended task

Anything else

In the tree view call modal, we are directly checking for the operator, which is probably the issue.

https://github.com/apache/airflow/blob/56a69b7c3e4fb3984b1e509ce2880c427a86e6ba/airflow/www/static/js/tree.js#L308

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project’s Code of Conduct

[uranusjr]

We can “fake” identity by serialising MRO information and implement __instancecheck__ and __subclasscheck__ on BaseOperator. But honestly my suggestion would be to avoid class introspection entirely. Instead of checking for the operator type, we should add an attribute (zoomable or something) and set it to True on SubDagOperator.

[bhavaniravi]

Done @bbovenzi