diff --git a/airflow-core/src/airflow/utils/dag_version_inflation_checker.py b/airflow-core/src/airflow/utils/dag_version_inflation_checker.py index 06fa96a7faae0..e9b415b255f1b 100644 --- a/airflow-core/src/airflow/utils/dag_version_inflation_checker.py +++ b/airflow-core/src/airflow/utils/dag_version_inflation_checker.py @@ -513,8 +513,6 @@ def visit_With(self, node: ast.With): for body in node.body: self.visit(body) - # Only exit Dag with block if we entered it; unconditional exit would prematurely - # reset the context when a non-DAG with-statement is nested inside a DAG with block. if is_with_dag_context: self.dag_detector.exit_dag_context() diff --git a/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py b/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py index 035375c01f348..91c468338ba6a 100644 --- a/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py +++ b/airflow-core/tests/unit/utils/test_dag_version_inflation_checker.py @@ -814,53 +814,25 @@ def test_task(): assert len(warnings) == 1 def test_nested_non_dag_with_does_not_exit_dag_context(self): - """Regression test: a non-DAG with-statement nested inside a DAG with-block must not - prematurely exit the DAG context. - - Before the fix, visit_With() called exit_dag_context() unconditionally regardless of - whether a DAG context was entered. This meant any nested non-DAG with-statement (e.g., - ``with open(...) as f``) would reset is_in_dag_context=False, causing tasks that appear - AFTER the nested with-statement (but still inside the DAG with-block) to be missed. - """ + """A non-Dag with-statement nested inside a Dag with-block must not exit the Dag context.""" code = """ from airflow import DAG from datetime import datetime from airflow.operators.bash import BashOperator with DAG('my_dag') as dag: - with open('some_file.txt') as f: # non-DAG with — must not exit DAG context + with open('some_file.txt') as f: data = f.read() - # This task is still inside the DAG block and uses datetime.now() — must be flagged t1 = BashOperator( task_id='test', - bash_command=str(datetime.now()), + bash_command=str(datetime.now()), # !problem ) -""" - warnings = self._check_code(code) - assert len(warnings) == 1, ( - "BashOperator after a nested non-DAG with-statement must still be detected " - "as inside the DAG context and flagged for using datetime.now()" - ) - - def test_multiple_nested_non_dag_withs_do_not_exit_dag_context(self): - """Multiple successive non-DAG with-statements inside a DAG block must not exit the context.""" - code = """ -from airflow import DAG -from datetime import datetime -from airflow.operators.bash import BashOperator - -with DAG('my_dag') as dag: - with open('a.txt') as f1: - pass - with open('b.txt') as f2: - pass - t1 = BashOperator(task_id='t', bash_command=str(datetime.now())) """ warnings = self._check_code(code) assert len(warnings) == 1 def test_task_inside_nested_non_dag_with_is_still_flagged(self): - """A task constructed inside a nested non-DAG with-block is still in DAG context.""" + """A task constructed inside a nested non-Dag with-block is still in Dag context.""" code = """ from airflow import DAG from datetime import datetime @@ -868,13 +840,13 @@ def test_task_inside_nested_non_dag_with_is_still_flagged(self): with DAG('my_dag') as dag: with open('a.txt') as f: - t1 = BashOperator(task_id='t', bash_command=str(datetime.now())) + t1 = BashOperator(task_id='t', bash_command=str(datetime.now())) # !problem """ warnings = self._check_code(code) assert len(warnings) == 1 - def test_task_outside_dag_with_not_flagged_after_nested_fix(self): - """Tasks genuinely outside any DAG block must not be flagged — the fix must not over-flag.""" + def test_task_outside_dag_with_not_flagged(self): + """Tasks genuinely outside any Dag block must not be flagged.""" code = """ from airflow import DAG from datetime import datetime @@ -884,8 +856,7 @@ def test_task_outside_dag_with_not_flagged_after_nested_fix(self): with open('a.txt') as f: pass -# This task is outside the DAG with block — should NOT be flagged t_outside = BashOperator(task_id='outside', bash_command=str(datetime.now())) """ warnings = self._check_code(code) - assert len(warnings) == 0, "A task constructed outside the DAG with-block must not be flagged" + assert len(warnings) == 0