From 4fe52796e0fa41556546eac25d7ce41785237574 Mon Sep 17 00:00:00 2001 From: GPK Date: Tue, 12 Nov 2024 12:17:20 +0000 Subject: [PATCH] Redirect old location module imports to standard provider (#43610) --------- Co-authored-by: Amogh Desai --- airflow/__init__.py | 38 ++++++++++++++ .../standard/test_module_redirect_finder.py | 51 +++++++++++++++++++ 2 files changed, 89 insertions(+) create mode 100644 providers/tests/standard/test_module_redirect_finder.py diff --git a/airflow/__init__.py b/airflow/__init__.py index 411aac70fc6f..a4a3bfbbfe7e 100644 --- a/airflow/__init__.py +++ b/airflow/__init__.py @@ -30,6 +30,7 @@ import os import sys import warnings +from importlib.abc import MetaPathFinder from typing import TYPE_CHECKING if os.environ.get("_AIRFLOW_PATCH_GEVENT"): @@ -133,3 +134,40 @@ def __getattr__(name: str): from airflow import plugins_manager plugins_manager.ensure_plugins_loaded() + + +class ModuleRedirectFinder(MetaPathFinder): + def __init__(self, redirects: dict[str, str]): + self.redirects: dict[str, str] = redirects + + def find_spec(self, old_module: str, path, target=None): + import importlib + + if old_module in self.redirects: + new_module_name = self.redirects[old_module] + warnings.warn( + f"Module '{old_module}' is deprecated, Please import it from '{new_module_name}'.", + DeprecationWarning, + stacklevel=2, + ) + return importlib.util.find_spec(new_module_name) + return None + + +module_redirects_references: dict[str, str] = { + "airflow.operators.python": "airflow.providers.standard.operators.python", + "airflow.operators.bash": "airflow.providers.standard.operators.bash", + "airflow.operators.datetime": "airflow.providers.standard.operators.datetime", + "airflow.operators.weekday": "airflow.providers.standard.operators.weekday", + "airflow.sensors.bash": "airflow.providers.standard.sensors.bash", + "airflow.sensors.date_time": "airflow.providers.standard.sensors.date_time", + "airflow.sensors.python": "airflow.providers.standard.sensors.python", + "airflow.sensors.time": "airflow.providers.standard.sensors.time", + "airflow.sensors.time_delta": "airflow.providers.standard.sensors.time_delta", + "airflow.sensors.weekday": "airflow.providers.standard.sensors.weekday", + "airflow.hooks.filesystem": "airflow.providers.standard.hooks.filesystem", + "airflow.hooks.package_index": "airflow.providers.standard.hooks.package_index", + "airflow.hooks.subprocess": "airflow.providers.standard.hooks.subprocess", + "airflow.utils.python_virtualenv": "airflow.providers.standard.utils.python_virtualenv", +} +sys.meta_path.append(ModuleRedirectFinder(module_redirects_references)) diff --git a/providers/tests/standard/test_module_redirect_finder.py b/providers/tests/standard/test_module_redirect_finder.py new file mode 100644 index 000000000000..0b95343f1ca0 --- /dev/null +++ b/providers/tests/standard/test_module_redirect_finder.py @@ -0,0 +1,51 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import importlib + +import pytest + +from tests_common.test_utils.compat import AIRFLOW_V_3_0_PLUS + + +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Valid for Airflow 3.0+ only") +class TestModuleRedirectFinder: + @pytest.mark.parametrize( + "old_module, new_module", + [ + ("airflow.operators.python", "airflow.providers.standard.operators.python"), + ("airflow.operators.bash", "airflow.providers.standard.operators.bash"), + ("airflow.operators.datetime", "airflow.providers.standard.operators.datetime"), + ("airflow.operators.weekday", "airflow.providers.standard.operators.weekday"), + ("airflow.sensors.bash", "airflow.providers.standard.sensors.bash"), + ("airflow.sensors.date_time", "airflow.providers.standard.sensors.date_time"), + ("airflow.sensors.python", "airflow.providers.standard.sensors.python"), + ("airflow.sensors.time", "airflow.providers.standard.sensors.time"), + ("airflow.sensors.time_delta", "airflow.providers.standard.sensors.time_delta"), + ("airflow.sensors.weekday", "airflow.providers.standard.sensors.weekday"), + ("airflow.hooks.filesystem", "airflow.providers.standard.hooks.filesystem"), + ("airflow.hooks.package_index", "airflow.providers.standard.hooks.package_index"), + ("airflow.hooks.subprocess", "airflow.providers.standard.hooks.subprocess"), + ("airflow.utils.python_virtualenv", "airflow.providers.standard.utils.python_virtualenv"), + ], + ) + def test_module_redirect_finder_for_old_location(self, old_module, new_module): + try: + assert importlib.import_module(old_module).__name__ == new_module + except ImportError: + pytest.fail(f"Failed to import module, incorrect reference: {old_module} -> {new_module}")