From dc93d1d545fbe4ca1f2284c0e1e8e2255b987aaf Mon Sep 17 00:00:00 2001 From: mohsinsurani Date: Sun, 7 Jun 2026 17:35:06 +0100 Subject: [PATCH 1/2] Add Dag file context to DagBag warnings --- .../src/airflow/dag_processing/dagbag.py | 19 ++++++++-- .../airflow/dag_processing/importers/base.py | 2 + .../importers/python_importer.py | 2 + .../tests/unit/dag_processing/test_dagbag.py | 38 +++++++++++++++++++ 4 files changed, 58 insertions(+), 3 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index dcaab93f1ab25..d92e6f44cb61e 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -75,6 +75,17 @@ def _capture_with_reraise() -> Generator[list[warnings.WarningMessage], None, No ) +def _format_dag_import_warning_message(message: str, source_path: str, dag_file_path: str | None) -> str: + """Include the imported Dag file when a warning points to another source file.""" + if ( + dag_file_path + and source_path != dag_file_path + and not source_path.startswith(f"{dag_file_path}{os.sep}") + ): + return f"{message} (Dag file: {dag_file_path})" + return message + + class FileLoadStat(NamedTuple): """ Information about single file. @@ -328,14 +339,16 @@ def process_file(self, filepath, only_if_updated=True, safe_mode=True): if result.warnings: formatted_warnings = [ - f"{w.file_path}:{w.line_number}: {w.warning_type}: {w.message}" for w in result.warnings + f"{w.file_path}:{w.line_number}: {w.warning_type}: " + f"{_format_dag_import_warning_message(w.message, w.file_path, w.dag_file_path)}" + for w in result.warnings ] self.captured_warnings[filepath] = tuple(formatted_warnings) # Re-emit warnings so they can be handled by Python's warning system for w in result.warnings: warnings.warn_explicit( - message=w.message, - category=UserWarning, + message=_format_dag_import_warning_message(w.message, w.file_path, w.dag_file_path), + category=w.category, filename=w.file_path, lineno=w.line_number or 0, ) diff --git a/airflow-core/src/airflow/dag_processing/importers/base.py b/airflow-core/src/airflow/dag_processing/importers/base.py index 643c16611f9b9..1cd175c8e1d1f 100644 --- a/airflow-core/src/airflow/dag_processing/importers/base.py +++ b/airflow-core/src/airflow/dag_processing/importers/base.py @@ -74,6 +74,8 @@ class DagImportWarning: message: str warning_type: str = "general" line_number: int | None = None + category: type[Warning] = UserWarning + dag_file_path: str | None = None @dataclass diff --git a/airflow-core/src/airflow/dag_processing/importers/python_importer.py b/airflow-core/src/airflow/dag_processing/importers/python_importer.py index 47f77a86f73ec..f0ed09ab9e0aa 100644 --- a/airflow-core/src/airflow/dag_processing/importers/python_importer.py +++ b/airflow-core/src/airflow/dag_processing/importers/python_importer.py @@ -187,6 +187,8 @@ def import_file( message=str(warn_msg.message), warning_type=category, line_number=warn_msg.lineno, + category=warn_msg.category, + dag_file_path=relative_path, ) ) diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index a673f9f1b0d2a..c42e0dd57938d 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1109,6 +1109,44 @@ def test_dabgag_captured_warnings(self): assert dag_file not in dagbag.captured_warnings assert dagbag.dagbag_stats[0].warning_num == 0 + def test_dagbag_warning_includes_dag_file_when_warning_source_differs(self, tmp_path): + dag_file = tmp_path / "test_dag_warning_from_internal_source.py" + warning_source = "/provider/deprecated_operator.py" + dag_file.write_text( + textwrap.dedent( + f"""\ + import warnings + from datetime import datetime + + from airflow.sdk import DAG + + warnings.warn_explicit( + "Deprecated provider behavior", + category=DeprecationWarning, + filename={warning_source!r}, + lineno=7, + ) + + with DAG("test_dag_warning_from_internal_source", start_date=datetime(2024, 1, 1)): + pass + """ + ) + ) + dagbag = DagBag(dag_folder=dag_file.as_posix(), include_examples=False, collect_dags=False) + + with pytest.warns( + DeprecationWarning, + match=f"Deprecated provider behavior .*Dag file: {re.escape(dag_file.as_posix())}", + ): + dagbag.collect_dags(dag_folder=dagbag.dag_folder, include_examples=False, only_if_updated=False) + + assert dagbag.captured_warnings == { + dag_file.as_posix(): ( + f"{warning_source}:7: DeprecationWarning: " + f"Deprecated provider behavior (Dag file: {dag_file.as_posix()})", + ) + } + @pytest.fixture def warning_zipped_dag_path(self, tmp_path: pathlib.Path) -> str: warnings_dag_file = TEST_DAGS_FOLDER / "test_dag_warnings.py" From 11f53d2f4d48f52c5ffbb964fb36b916e49aabbf Mon Sep 17 00:00:00 2001 From: mohsinsurani Date: Mon, 17 Aug 2026 13:38:41 +0530 Subject: [PATCH 2/2] Improve Dag warning source coverage --- .../src/airflow/dag_processing/dagbag.py | 8 ++------ .../airflow/dag_processing/importers/base.py | 2 +- .../tests/unit/dag_processing/test_dagbag.py | 20 +++++++++---------- 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 7204aaf4be333..64f66d556e3a8 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -75,13 +75,9 @@ def _capture_with_reraise() -> Generator[list[warnings.WarningMessage], None, No ) -def _format_dag_import_warning_message(message: str, source_path: str, dag_file_path: str | None) -> str: +def _format_dag_import_warning_message(message: str, source_path: str, dag_file_path: str) -> str: """Include the imported Dag file when a warning points to another source file.""" - if ( - dag_file_path - and source_path != dag_file_path - and not source_path.startswith(f"{dag_file_path}{os.sep}") - ): + if source_path != dag_file_path and not source_path.startswith(f"{dag_file_path}{os.sep}"): return f"{message} (Dag file: {dag_file_path})" return message diff --git a/airflow-core/src/airflow/dag_processing/importers/base.py b/airflow-core/src/airflow/dag_processing/importers/base.py index 1cd175c8e1d1f..7067e6472b56c 100644 --- a/airflow-core/src/airflow/dag_processing/importers/base.py +++ b/airflow-core/src/airflow/dag_processing/importers/base.py @@ -75,7 +75,7 @@ class DagImportWarning: warning_type: str = "general" line_number: int | None = None category: type[Warning] = UserWarning - dag_file_path: str | None = None + dag_file_path: str = "" @dataclass diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 0f41df38450bb..1c045d070ddd5 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1146,9 +1146,12 @@ def test_dabgag_captured_warnings(self): assert dag_file not in dagbag.captured_warnings assert dagbag.dagbag_stats[0].warning_num == 0 - def test_dagbag_warning_includes_dag_file_when_warning_source_differs(self, tmp_path): + @pytest.mark.parametrize("warning_from_dag_file", [True, False]) + def test_dagbag_warning_includes_dag_file_when_warning_source_differs( + self, tmp_path, warning_from_dag_file + ): dag_file = tmp_path / "test_dag_warning_from_internal_source.py" - warning_source = "/provider/deprecated_operator.py" + warning_source = dag_file.as_posix() if warning_from_dag_file else "/provider/deprecated_operator.py" dag_file.write_text( textwrap.dedent( f"""\ @@ -1170,18 +1173,15 @@ def test_dagbag_warning_includes_dag_file_when_warning_source_differs(self, tmp_ ) ) dagbag = DagBag(dag_folder=dag_file.as_posix(), include_examples=False, collect_dags=False) + expected_message = "Deprecated provider behavior" + if not warning_from_dag_file: + expected_message += f" (Dag file: {dag_file.as_posix()})" - with pytest.warns( - DeprecationWarning, - match=f"Deprecated provider behavior .*Dag file: {re.escape(dag_file.as_posix())}", - ): + with pytest.warns(DeprecationWarning, match=re.escape(expected_message)): dagbag.collect_dags(dag_folder=dagbag.dag_folder, include_examples=False, only_if_updated=False) assert dagbag.captured_warnings == { - dag_file.as_posix(): ( - f"{warning_source}:7: DeprecationWarning: " - f"Deprecated provider behavior (Dag file: {dag_file.as_posix()})", - ) + dag_file.as_posix(): (f"{warning_source}:7: DeprecationWarning: {expected_message}",) } @pytest.fixture