[AIRFLOW-3060] DAG context manager fails to exit properly in certain circumstances

Author: newtonle

airflow/models.py

@@ -3442,6 +3442,8 @@ def __init__(
         self.on_success_callback = on_success_callback
         self.on_failure_callback = on_failure_callback
 
+        self._context_manager_set = False
+
         self._comps = {
             'dag_id',
             'task_ids',
@@ -3490,13 +3492,16 @@ def __hash__(self):
 
     def __enter__(self):
         global _CONTEXT_MANAGER_DAG
-        self._old_context_manager_dag = _CONTEXT_MANAGER_DAG
-        _CONTEXT_MANAGER_DAG = self
+        if not self._context_manager_set:
+            self._old_context_manager_dag = _CONTEXT_MANAGER_DAG
+            _CONTEXT_MANAGER_DAG = self
+            self._context_manager_set = True
         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 ----------------------------------------------
 

tests/models.py

@@ -147,6 +147,15 @@ def test_dag_as_context_manager(self):
         self.assertEqual(dag.dag_id, 'creating_dag_in_cm')
         self.assertEqual(dag.tasks[0].task_id, 'op6')
 
+        with dag:
+            with dag:
+                op7 = DummyOperator(task_id='op7')
+        op8 = DummyOperator(task_id='op8')
+        op8.dag = dag2
+
+        self.assertEqual(op7.dag, dag)
+        self.assertEqual(op8.dag, dag2)
+
     def test_dag_topological_sort(self):
         dag = DAG(
             'dag',