[AIRFLOW-881] Check if SubDagOperator is in DAG context manager

When initializing a SubDagOperator, the `dag`
param should not be
required if it is within a DAG context manager. So
we check if that
is the case and use that as the parent DAG if
found (and `dag` param
is not specified).

Closes apache#2087 from dhuang/AIRFLOW-881

airflow/operators/subdag_operator.py

@@ -41,9 +41,11 @@ def __init__(
         :param dag: the parent DAG
         :type subdag: airflow.DAG
         """
-        if 'dag' not in kwargs:
-            raise AirflowException("Please pass in the `dag` param")
-        dag = kwargs['dag']
+        import airflow.models
+        dag = kwargs.get('dag') or airflow.models._CONTEXT_MANAGER_DAG
+        if not dag:
+            raise AirflowException('Please pass in the `dag` param or call '
+                                   'within a DAG context manager')
         session = kwargs.pop('session')
         super(SubDagOperator, self).__init__(*args, **kwargs)
 

tests/operators/subdag_operator.py

@@ -54,6 +54,17 @@ def test_subdag_name(self):
             AirflowException,
             SubDagOperator, task_id='test', dag=dag, subdag=subdag_bad3)
 
+    def test_subdag_in_context_manager(self):
+        """
+        Creating a sub DAG within a main DAG's context manager
+        """
+        with DAG('parent', default_args=default_args) as dag:
+            subdag = DAG('parent.test', default_args=default_args)
+            op = SubDagOperator(task_id='test', subdag=subdag)
+
+            self.assertEqual(op.dag, dag)
+            self.assertEqual(op.subdag, subdag)
+
     def test_subdag_pools(self):
         """
         Subdags and subdag tasks can't both have a pool with 1 slot