Replace DeprecationWarning with DeprecatedImportWarning in airflow-core (apache#56836)

Author: KoviAnusha

airflow-core/src/airflow/__init__.py

@@ -33,6 +33,8 @@
 import warnings
 from typing import TYPE_CHECKING
 
+from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
 if os.environ.get("_AIRFLOW_PATCH_GEVENT"):
     # If you are using gevents and start airflow webserver, you might want to run gevent monkeypatching
     # as one of the first thing when Airflow is started. This allows gevent to patch networking and other
@@ -105,7 +107,7 @@ def __getattr__(name: str):
         warnings.warn(
             f"Import {name!r} directly from the airflow module is deprecated and "
             f"will be removed in the future. Please import it from 'airflow{module_path}.{attr_name}'.",
-            DeprecationWarning,
+            DeprecatedImportWarning,
             stacklevel=2,
         )
 

airflow-core/src/airflow/datasets/__init__.py

@@ -27,6 +27,8 @@
 import importlib
 import warnings
 
+from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
 # TODO: Remove this module in Airflow 3.2
 
 _names_moved = {
@@ -47,7 +49,7 @@ def __getattr__(name: str):
     warnings.warn(
         f"Import 'airflow.datasets.{name}' is deprecated and "
         f"will be removed in Airflow 3.2. Please import it from '{module_path}.{new_name}'.",
-        DeprecationWarning,
+        DeprecatedImportWarning,
         stacklevel=2,
     )
     mod = importlib.import_module(module_path, __name__)

airflow-core/src/airflow/datasets/metadata.py

@@ -20,13 +20,14 @@
 import warnings
 
 from airflow.sdk import Metadata
+from airflow.utils.deprecation_tools import DeprecatedImportWarning
 
 # TODO: Remove this module in Airflow 3.2
 
 warnings.warn(
     "Import from the airflow.datasets.metadata module is deprecated and will "
     "be removed in Airflow 3.2. Please import it from 'airflow.sdk'.",
-    DeprecationWarning,
+    DeprecatedImportWarning,
     stacklevel=2,
 )
 

airflow-core/src/airflow/models/dag.py

@@ -750,10 +750,12 @@ def __getattr__(name: str):
 
     import warnings
 
+    from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
     warnings.warn(
         f"Import {name!r} directly from the airflow module is deprecated and "
         f"will be removed in the future. Please import it from 'airflow.sdk'.",
-        DeprecationWarning,
+        DeprecatedImportWarning,
         stacklevel=2,
     )
 

airflow-core/src/airflow/models/dagbag.py

@@ -146,10 +146,12 @@ def __getattr__(name: str) -> Any:
     if name in {"DagBag", "FileLoadStat", "timeout"}:
         import warnings
 
+        from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
         warnings.warn(
             f"Importing {name} from airflow.models.dagbag is deprecated and will be removed in a future "
             "release. Please import from airflow.dag_processing.dagbag instead.",
-            DeprecationWarning,
+            DeprecatedImportWarning,
             stacklevel=2,
         )
         # Import on demand to avoid import-time side effects

airflow-core/src/airflow/secrets/__init__.py

@@ -51,10 +51,12 @@ def __getattr__(name):
     if name == "DEFAULT_SECRETS_SEARCH_PATH_WORKERS":
         import warnings
 
+        from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
         warnings.warn(
             "airflow.secrets.DEFAULT_SECRETS_SEARCH_PATH_WORKERS is moved to the Task SDK. "
             "Use airflow.sdk.execution_time.secrets.DEFAULT_SECRETS_SEARCH_PATH_WORKERS instead.",
-            DeprecationWarning,
+            DeprecatedImportWarning,
             stacklevel=2,
         )
         try:

airflow-core/src/airflow/timetables/datasets.py

@@ -19,6 +19,7 @@
 import warnings
 
 from airflow.timetables.assets import AssetOrTimeSchedule
+from airflow.utils.deprecation_tools import DeprecatedImportWarning
 
 
 class DatasetOrTimeSchedule:
@@ -27,7 +28,7 @@ class DatasetOrTimeSchedule:
     def __new__(cls, *, timetable, datasets) -> AssetOrTimeSchedule:  # type: ignore[misc]
         warnings.warn(
             "DatasetOrTimeSchedule is deprecated and will be removed in Airflow 3.2. Use `airflow.timetables.AssetOrTimeSchedule` instead.",
-            DeprecationWarning,
+            DeprecatedImportWarning,
             stacklevel=2,
         )
         return AssetOrTimeSchedule(timetable=timetable, assets=datasets)

airflow-core/src/airflow/utils/dag_parsing_context.py

@@ -21,13 +21,14 @@
 
 from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
 from airflow.sdk.definitions.context import get_parsing_context
+from airflow.utils.deprecation_tools import DeprecatedImportWarning
 
 # TODO: Remove this module in Airflow 3.2
 
 warnings.warn(
     "Import from the airflow.utils.dag_parsing_context module is deprecated and "
     "will be removed in Airflow 3.2. Please import it from 'airflow.sdk'.",
-    DeprecationWarning,
+    DeprecatedImportWarning,
     stacklevel=2,
 )
 

airflow-core/tests/unit/core/test_airflow_module.py

@@ -23,7 +23,9 @@
 
 
 def test_deprecated_exception():
+    from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
     warning_pattern = "Import 'AirflowException' directly from the airflow module is deprecated"
-    with pytest.warns(DeprecationWarning, match=warning_pattern):
+    with pytest.warns(DeprecatedImportWarning, match=warning_pattern):
         # If there is no warning, then most possible it imported somewhere else.
         assert getattr(airflow, "AirflowException") is AirflowException

airflow-core/tests/unit/datasets/test_dataset.py

@@ -84,7 +84,9 @@ def test_backward_compat_import_before_airflow_3_2(module_path, attr_name, expec
         mod = importlib.import_module(module_path, __name__)
         attr = getattr(mod, attr_name)
     assert f"{attr.__module__}.{attr.__name__}" == expected_value
-    assert record[0].category is DeprecationWarning
+    from airflow.utils.deprecation_tools import DeprecatedImportWarning
+
+    assert record[0].category is DeprecatedImportWarning
     assert str(record[0].message) == warning_message