Fix Version Check for CLI Imports in Celery provider (#45255)

providers/src/airflow/providers/celery/cli/celery_command.py

@@ -33,6 +33,7 @@
 
 from airflow import settings
 from airflow.configuration import conf
+from airflow.providers.celery.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils import cli as cli_utils
 from airflow.utils.cli import setup_locations
 from airflow.utils.serve_logs import serve_logs
@@ -42,8 +43,10 @@
 
 def _run_command_with_daemon_option(*args, **kwargs):
     try:
-        from airflow.cli.commands.local_commands.daemon_utils import run_command_with_daemon_option
-
+        if AIRFLOW_V_3_0_PLUS:
+            from airflow.cli.commands.local_commands.daemon_utils import run_command_with_daemon_option
+        else:
+            from airflow.cli.commands.daemon_utils import run_command_with_daemon_option
         run_command_with_daemon_option(*args, **kwargs)
     except ImportError:
         from airflow.exceptions import AirflowOptionalProviderFeatureException

providers/src/airflow/providers/celery/executors/celery_executor.py

@@ -37,9 +37,7 @@
 
 from celery import states as celery_states
 from deprecated import deprecated
-from packaging.version import Version
 
-from airflow import __version__ as airflow_version
 from airflow.cli.cli_config import (
     ARG_DAEMON,
     ARG_LOG_FILE,
@@ -56,6 +54,7 @@
 from airflow.configuration import conf
 from airflow.exceptions import AirflowProviderDeprecationWarning, AirflowTaskTimeout
 from airflow.executors.base_executor import BaseExecutor
+from airflow.providers.celery.version_compat import AIRFLOW_V_2_8_PLUS
 from airflow.stats import Stats
 from airflow.utils.state import TaskInstanceState
 
@@ -161,11 +160,10 @@ def __getattr__(name):
     action="store_true",
 )
 
-CELERY_CLI_COMMAND_PATH = (
-    "airflow.providers.celery.cli.celery_command"
-    if Version(airflow_version) >= Version("2.8.0")
-    else "airflow.cli.commands.local_commands.celery_command"
-)
+if AIRFLOW_V_2_8_PLUS:
+    CELERY_CLI_COMMAND_PATH = "airflow.providers.celery.cli.celery_command"
+else:
+    CELERY_CLI_COMMAND_PATH = "airflow.cli.commands.celery_command"
 
 CELERY_COMMANDS = (
     ActionCommand(

providers/src/airflow/providers/celery/version_compat.py

@@ -0,0 +1,30 @@
+# 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
+
+
+def get_base_airflow_version_tuple() -> tuple[int, int, int]:
+    from packaging.version import Version
+
+    from airflow import __version__
+
+    airflow_version = Version(__version__)
+    return airflow_version.major, airflow_version.minor, airflow_version.micro
+
+
+AIRFLOW_V_2_8_PLUS = get_base_airflow_version_tuple() >= (2, 8, 0)
+AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)

providers/tests/celery/cli/test_celery_command.py

@@ -32,7 +32,7 @@
 from airflow.providers.celery.cli import celery_command
 
 from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_2_10_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_2_10_PLUS, AIRFLOW_V_3_0_PLUS
 
 pytestmark = pytest.mark.db_test
 
@@ -270,11 +270,7 @@ def test_run_command(self, mock_celery_app):
             ]
         )
 
-    @mock.patch("airflow.cli.commands.local_commands.daemon_utils.TimeoutPIDLockFile")
-    @mock.patch("airflow.cli.commands.local_commands.daemon_utils.setup_locations")
-    @mock.patch("airflow.cli.commands.local_commands.daemon_utils.daemon")
-    @mock.patch("airflow.providers.celery.executors.celery_executor.app")
-    def test_run_command_daemon(self, mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file):
+    def _test_run_command_daemon(self, mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file):
         mock_setup_locations.return_value = (
             mock.MagicMock(name="pidfile"),
             mock.MagicMock(name="stdout"),
@@ -363,3 +359,23 @@ def test_run_command_daemon(self, mock_celery_app, mock_daemon, mock_setup_locat
             mock.call().__exit__(None, None, None),
             mock.call().__exit__(None, None, None),
         ]
+
+    @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow 3.0-")
+    @mock.patch("airflow.cli.commands.daemon_utils.TimeoutPIDLockFile")
+    @mock.patch("airflow.cli.commands.daemon_utils.setup_locations")
+    @mock.patch("airflow.cli.commands.daemon_utils.daemon")
+    @mock.patch("airflow.providers.celery.executors.celery_executor.app")
+    def test_run_command_daemon_v_3_below(
+        self, mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file
+    ):
+        self._test_run_command_daemon(mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file)
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Test requires Airflow 3.0+")
+    @mock.patch("airflow.cli.commands.local_commands.daemon_utils.TimeoutPIDLockFile")
+    @mock.patch("airflow.cli.commands.local_commands.daemon_utils.setup_locations")
+    @mock.patch("airflow.cli.commands.local_commands.daemon_utils.daemon")
+    @mock.patch("airflow.providers.celery.executors.celery_executor.app")
+    def test_run_command_daemon_v3_above(
+        self, mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file
+    ):
+        self._test_run_command_daemon(mock_celery_app, mock_daemon, mock_setup_locations, mock_pid_file)