[AIRFLOW-3123] Use a stack for DAG context management (apache#3956)

Author: newtonle

airflow/models.py

@@ -3442,7 +3442,7 @@ def __init__(
         self.on_success_callback = on_success_callback
         self.on_failure_callback = on_failure_callback
 
-        self._context_manager_set = False
+        self._old_context_manager_dags = []
 
         self._comps = {
             'dag_id',
@@ -3492,16 +3492,13 @@ def __hash__(self):
 
     def __enter__(self):
         global _CONTEXT_MANAGER_DAG
-        if not self._context_manager_set:
-            self._old_context_manager_dag = _CONTEXT_MANAGER_DAG
-            _CONTEXT_MANAGER_DAG = self
-            self._context_manager_set = True
+        self._old_context_manager_dags.append(_CONTEXT_MANAGER_DAG)
+        _CONTEXT_MANAGER_DAG = self
         return self
 
     def __exit__(self, _type, _value, _tb):
         global _CONTEXT_MANAGER_DAG
-        _CONTEXT_MANAGER_DAG = self._old_context_manager_dag
-        self._context_manager_set = False
+        _CONTEXT_MANAGER_DAG = self._old_context_manager_dags.pop()
 
     # /Context Manager ----------------------------------------------
 

tests/models.py

@@ -150,11 +150,13 @@ def test_dag_as_context_manager(self):
         with dag:
             with dag:
                 op7 = DummyOperator(task_id='op7')
-        op8 = DummyOperator(task_id='op8')
-        op8.dag = dag2
+            op8 = DummyOperator(task_id='op8')
+        op9 = DummyOperator(task_id='op8')
+        op9.dag = dag2
 
         self.assertEqual(op7.dag, dag)
-        self.assertEqual(op8.dag, dag2)
+        self.assertEqual(op8.dag, dag)
+        self.assertEqual(op9.dag, dag2)
 
     def test_dag_topological_sort(self):
         dag = DAG(