Airflow - generate DAG from custom object

Question

I'm using airflow 1.10.3. I am trying to create a class that I can instantiate multiple times to generate DAGs for me. Here is a simplified example:

from airflow import DAG
from airflow.operators.python_operator import PythonOperator


class DagClass:

    def __init__(self, dag_id):
        self.dag_id = dag_id


    def test_function():
        print("test 1")

    def test_function2():
        print("test 2")


    def generate_dag(self):
        dag = DAG(self.dag_id)

        task1 = PythonOperator(
                  task_id="test",
                  python_callable=self.test_function,
                  provide_context=True,
                  dag=dag
                )

        task2 = PythonOperator(
                task_id="test2",
                python_callable=self.test_function2,
                provide_context=True,
                dag=dag
                )

        task1.set_downstream(task2)

        return dag

Then, I try to use that class to create a DAG in another file here:

from airflow import DAG
from dynamic_dag_test.dag_class import DagClass

dag_class = DagClass('test_dynamic_dag')

dag = dag_class.generate_dag()

globals()['test_dynamic_dag'] = dag

However if I do that, I won't get any errors, but my DAG will never show up in the web UI. BUT, if I instead get rid of the class and just directly import the function from another file, it works fine.

Does anyone know why this is and how I can get it to work with the class? Thanks!

Answer 1

I can confirm the comment of David Schuler as well with Airflow 1.10.1. We had to put the import statement ( from airflow import DAG ) at the top of the Python file.