diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 584840f25c329..937a5b808df9f 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -76,6 +76,13 @@ 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) -> str: + """Include the imported Dag file when a warning points to another source file.""" + 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 + + class FileLoadStat(NamedTuple): """ Information about single file. @@ -345,14 +352,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..7067e6472b56c 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 = "" @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 cb110ef5c7fda..a22553a8287a3 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1214,6 +1214,44 @@ def test_dabgag_captured_warnings(self): assert dag_file not in dagbag.captured_warnings assert dagbag.dagbag_stats[0].warning_num == 0 + @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 = dag_file.as_posix() if warning_from_dag_file else "/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) + 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=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: {expected_message}",) + } + @pytest.fixture def warning_zipped_dag_path(self, tmp_path: pathlib.Path) -> str: warnings_dag_file = TEST_DAGS_FOLDER / "test_dag_warnings.py"