This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new ff8b7765b22 [v3-3-test] Detect Dag subclasses in the Dag version
inflation check (#72861) (#72898)
ff8b7765b22 is described below
commit ff8b7765b22daec0ebfdd5ff9cc8861913bb2121
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Fri Sep 11 14:25:17 2026 +0530
[v3-3-test] Detect Dag subclasses in the Dag version inflation check
(#72861) (#72898)
* [v3-3-test] Detect Dag subclasses in the Dag version inflation check
(#72861)
The check only recognized a Dag constructor when the callable resolved
through a from-import of airflow's DAG or dag. A Dag built from a
subclass, or reached through an aliased module import, produced no
warnings at all — and because task detection keys off the same
predicate, neither did any task inside its with-block.
The Dag file is not imported yet when the check runs, so there is no
object to test — only names and the file's own imports. The cost of
leaning on names is an occasional false positive on something merely
named like a Dag.
(cherry picked from commit 1e2ad803f74eadb5fb28d4504ca06b008c48b63e)
Co-authored-by: Jed Cunningham
<[email protected]>
* Add missing dedent import to fix the version inflation checker tests
---------
Co-authored-by: Jed Cunningham
<[email protected]>
Co-authored-by: Rahul Vats <[email protected]>
Co-authored-by: Rahul Vats <[email protected]>
---
.../airflow/utils/dag_version_inflation_checker.py | 40 +++++---
.../utils/test_dag_version_inflation_checker.py | 109 +++++++++++++++------
2 files changed, 103 insertions(+), 46 deletions(-)
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 e9b415b255f..0b9b97a3b67 100644
--- a/airflow-core/src/airflow/utils/dag_version_inflation_checker.py
+++ b/airflow-core/src/airflow/utils/dag_version_inflation_checker.py
@@ -312,24 +312,35 @@ class DagTaskDetector:
def is_dag_constructor(self, node: ast.Call) -> bool:
"""Check if a call is a Dag constructor."""
- # to handle use case "from airflow import sdk" and "with sdk.DAG()"
- if isinstance(node.func, ast.Attribute) and
isinstance(node.func.value, ast.Name):
- if node.func.value.id in self.from_imports:
- module, original = self.from_imports[node.func.value.id]
- if (module == "airflow" or module.startswith("airflow.")) and
node.func.attr in (
- "DAG",
- "dag",
- ):
+ # The Dag file is not imported yet, so there is no object to test —
only names and its imports.
+ # A *DAG/*Dag suffix counts on its own, so a subclass named anything
else is missed.
+ # Lowercase "dag" is an ordinary word and needs an import — and only a
plain name can
+ # be an imported one, since "from airflow import DAG as D" says
nothing about config.D().
+ func = node.func
+
+ # DAG(...), TeamDAG(...), an alias like D(...), or the @dag(...)
decorator
+ if isinstance(func, ast.Name):
+ if func.id in self.from_imports:
+ module, original = self.from_imports[func.id]
+ if self._is_airflow_module(module) and original in ("DAG",
"dag"):
return True
+ return func.id.endswith(("DAG", "Dag"))
- # to handle use case "from airflow import DAG" form or "from
airflow.decorator import dag"
- if isinstance(node.func, ast.Name) and node.func.id in
self.from_imports:
- module, original = self.from_imports[node.func.id]
- if (module == "airflow" or module.startswith("airflow.")) and
original in ("DAG", "dag"):
- return True
+ # sdk.DAG(...), airflow.sdk.DAG(...), or the @sdk.dag(...) decorator
+ if isinstance(func, ast.Attribute):
+ if func.attr == "dag" and isinstance(func.value, ast.Name):
+ owner = self.from_imports.get(func.value.id)
+ if owner and self._is_airflow_module(owner[0]):
+ return True
+ return func.attr.endswith(("DAG", "Dag"))
return False
+ @staticmethod
+ def _is_airflow_module(module: str) -> bool:
+ """Check whether a module path is "airflow" or anything under it."""
+ return module == "airflow" or module.startswith("airflow.")
+
def is_task_constructor(self, node: ast.Call) -> bool:
"""
Check if a call is a Task constructor.
@@ -507,13 +518,14 @@ class AirflowRuntimeVaryingValueChecker(ast.NodeVisitor):
if item.optional_vars and isinstance(item.optional_vars,
ast.Name):
self._register_dag_instances([item.optional_vars])
+ already_in_dag_context = self.dag_detector.is_in_dag_context
if is_with_dag_context:
self.dag_detector.enter_dag_context()
for body in node.body:
self.visit(body)
- if is_with_dag_context:
+ if is_with_dag_context and not already_in_dag_context:
self.dag_detector.exit_dag_context()
def visit_FunctionDef(self, node: ast.FunctionDef):
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 91c468338ba..992ca2352bc 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
@@ -17,6 +17,9 @@
from __future__ import annotations
import ast
+from textwrap import dedent
+
+import pytest
from airflow.utils.dag_version_inflation_checker import (
AirflowRuntimeVaryingValueChecker,
@@ -276,42 +279,40 @@ class TestDAGTaskDetector:
self.from_imports = {}
self.detector = DagTaskDetector(self.from_imports)
- def test_is_dag_constructor__detects_traditional_dag_call_uppercase(self):
- """
- Detect uppercase DAG() when imported.
-
- Usage: dag = DAG(dag_id="my_dag")
- """
- code = 'DAG(dag_id="my_dag")'
- call_node = ast.parse(code, mode="eval").body
- self.from_imports["DAG"] = ("airflow", "DAG")
-
- result = self.detector.is_dag_constructor(call_node)
-
- assert result is True
-
- def test_is_dag_constructor__detects_dag_generated_by_decorator(self):
- """
- Detect Dag generated by decorator.
-
- Usage: @dag(dag_id="my_dag")
- """
- code = 'dag(dag_id="my_dag")'
- call_node = ast.parse(code, mode="eval").body
- self.from_imports["dag"] = ("airflow.decorators", "dag")
-
- result = self.detector.is_dag_constructor(call_node)
-
- assert result is True
-
- def test_is_dag_constructor__ignores_non_dag_functions(self):
- """Regular function calls should not be detected as Dag
constructors."""
- code = "my_function()"
+ @pytest.mark.parametrize(
+ ("code", "from_imports", "expected"),
+ [
+ # Bound to airflow's DAG or dag by an import, under any alias.
+ ('DAG(dag_id="my_dag")', {"DAG": ("airflow", "DAG")}, True),
+ ('dag(dag_id="my_dag")', {"dag": ("airflow.decorators", "dag")},
True),
+ ("D(x=1)", {"D": ("airflow", "DAG")}, True),
+ # A *DAG/*Dag suffix is a class-naming convention, so it stands on
its own.
+ ('MyCustomDAG(dag_id="my_dag")', {}, True),
+ ('TeamDag(dag_id="my_dag")', {}, True),
+ ('sdk.DAG(dag_id="my_dag")', {}, True),
+ ('airflow.sdk.DAG(dag_id="my_dag")', {}, True),
+ # Lowercase is not a naming convention, so on an attribute it
needs an airflow owner.
+ ('sdk.dag(dag_id="x")', {"sdk": ("airflow", "sdk")}, True),
+ ('company.dag(dag_id="x")', {"company": ("company", "tools")},
False),
+ ('company.dag(dag_id="x")', {}, False),
+ # A bare lowercase name proves nothing at all, so the suffix stays
case-sensitive.
+ ("dag(x=1)", {}, False),
+ ("dag(x=1)", {"dag": ("company.lib", "dag")}, False),
+ ("create_dag()", {}, False),
+ ("get_dag()", {}, False),
+ ("my_function()", {}, False),
+ # An import binds a bare name, so an attribute segment is never an
alias.
+ ("config.D(x=1)", {"D": ("airflow", "DAG")}, False),
+ ],
+ )
+ def test_is_dag_constructor(self, code, from_imports, expected):
+ """Whether a call is treated as constructing a Dag, given the file's
imports."""
+ self.from_imports.update(from_imports)
call_node = ast.parse(code, mode="eval").body
result = self.detector.is_dag_constructor(call_node)
- assert result is False
+ assert result is expected
def test_is_task_constructor__true_when_inside_dag_context(self):
"""
@@ -782,6 +783,32 @@ with sdk.DAG(
warnings = self._check_code(code)
assert len(warnings) == 1
+ def test_dag_subclass_from_another_module(self):
+ code = dedent(
+ """
+ from company.dags import TeamDAG
+ from datetime import datetime
+
+ with TeamDAG(dag_id=f"team_{datetime.now()}") as dag:
+ pass
+ """
+ )
+ warnings = self._check_code(code)
+ assert len(warnings) == 1
+
+ def test_import_dag_from_aliased_sdk_module(self):
+ code = dedent(
+ """
+ import airflow.sdk as sdk
+ import datetime
+
+ with sdk.DAG(dag_id=f"test_{datetime.datetime.now()}") as dag:
+ pass
+ """
+ )
+ warnings = self._check_code(code)
+ assert len(warnings) == 1
+
def test_python_task_with_task_decorator(self):
code = """
from airflow.sdk import task, DAG
@@ -831,6 +858,24 @@ with DAG('my_dag') as dag:
warnings = self._check_code(code)
assert len(warnings) == 1
+ def test_nested_dag_named_with_does_not_exit_outer_dag_context(self):
+ """A nested with-statement whose context manager merely looks like a
Dag must not exit the Dag context."""
+ code = dedent(
+ """
+ from airflow import DAG
+ from company.tools import ConfigDag
+ from datetime import datetime
+ from airflow.providers.standard.operators.bash import BashOperator
+
+ with DAG('my_dag') as dag:
+ with ConfigDag('x'):
+ pass
+ t1 = BashOperator(task_id='t',
bash_command=str(datetime.now())) # !problem
+ """
+ )
+ 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."""
code = """
