(airflow) branch main updated: Speedup determine installed airflow version in `ExternalPythonOperator` (#37409)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 8c0a4de663 Speedup determine installed airflow version in `ExternalPythonOperator` (#37409)
8c0a4de663 is described below

commit 8c0a4de66344647c759b6188cabafacf26ace9e3
Author: Andrey Anshin <An...@taragol.is>
AuthorDate: Wed Feb 14 05:16:58 2024 +0400

    Speedup determine installed airflow version in `ExternalPythonOperator` (#37409)
---
 airflow/operators/python.py    | 33 ++++++++++++++++++++++++-----
 tests/operators/test_python.py | 47 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 74 insertions(+), 6 deletions(-)

diff --git a/airflow/operators/python.py b/airflow/operators/python.py
index 0f005f43b2..2fa004f028 100644
--- a/airflow/operators/python.py
+++ b/airflow/operators/python.py
@@ -889,21 +889,44 @@ class ExternalPythonOperator(_BasePythonVirtualenvOperator):
                 )
             return False
 
+    @property
+    def _external_airflow_version_script(self):
+        """
+        Return python script which determines the version of the Apache Airflow.
+
+        Import airflow as a module might take a while as a result,
+        obtaining a version would take up to 1 second.
+        On the other hand, `importlib.metadata.version` will retrieve the package version pretty fast
+        something below 100ms; this includes new subprocess overhead.
+
+        Possible side effect: it might be a situation that backport package is not available
+        in Python 3.8 and below, which indicates that venv doesn't contain an `apache-airflow`
+        or something wrong with the environment.
+        """
+        return textwrap.dedent(
+            """
+            import sys
+            if sys.version_info >= (3, 9):
+                from importlib.metadata import version
+            else:
+                from importlib_metadata import version
+            print(version("apache-airflow"))
+            """
+        )
+
     def _get_airflow_version_from_target_env(self) -> str | None:
         from airflow import __version__ as airflow_version
 
         try:
             result = subprocess.check_output(
-                [self.python, "-c", "from airflow import __version__; print(__version__)"],
+                [self.python, "-c", self._external_airflow_version_script],
                 text=True,
-                # Avoid Airflow logs polluting stdout.
-                env={**os.environ, "_AIRFLOW__AS_LIBRARY": "true"},
             )
             target_airflow_version = result.strip()
             if target_airflow_version != airflow_version:
                 raise AirflowConfigException(
-                    f"The version of Airflow installed for the {self.python}("
-                    f"{target_airflow_version}) is different than the runtime Airflow version: "
+                    f"The version of Airflow installed for the {self.python} "
+                    f"({target_airflow_version}) is different than the runtime Airflow version: "
                     f"{airflow_version}. Make sure your environment has the same Airflow version "
                     f"installed as the Airflow runtime."
                 )
diff --git a/tests/operators/test_python.py b/tests/operators/test_python.py
index ba008ee1ca..f1a297aab3 100644
--- a/tests/operators/test_python.py
+++ b/tests/operators/test_python.py
@@ -1208,9 +1208,10 @@ class TestExternalPythonOperator(BaseTestPythonVirtualenvOperator):
         def f():
             return 1
 
-        task = PythonVirtualenvOperator(
+        task = ExternalPythonOperator(
             python_callable=f,
             task_id="task",
+            python=sys.executable,
             dag=self.dag,
         )
 
@@ -1218,6 +1219,50 @@ class TestExternalPythonOperator(BaseTestPythonVirtualenvOperator):
         with pytest.raises(DeserializingResultError):
             task._read_result(path=mock.Mock())
 
+    def test_airflow_version(self):
+        def f():
+            return 42
+
+        op = ExternalPythonOperator(
+            python_callable=f, task_id="task", python=sys.executable, expect_airflow=True
+        )
+        assert op._get_airflow_version_from_target_env()
+
+    def test_airflow_version_doesnt_match(self, caplog):
+        def f():
+            return 42
+
+        op = ExternalPythonOperator(
+            python_callable=f, task_id="task", python=sys.executable, expect_airflow=True
+        )
+
+        with mock.patch.object(
+            ExternalPythonOperator, "_external_airflow_version_script", new_callable=mock.PropertyMock
+        ) as mock_script:
+            mock_script.return_value = "print('1.10.4')"
+            caplog.set_level("WARNING")
+            caplog.clear()
+            assert op._get_airflow_version_from_target_env() is None
+            assert "(1.10.4) is different than the runtime Airflow version" in caplog.text
+
+    def test_airflow_version_script_error_handle(self, caplog):
+        def f():
+            return 42
+
+        op = ExternalPythonOperator(
+            python_callable=f, task_id="task", python=sys.executable, expect_airflow=True
+        )
+
+        with mock.patch.object(
+            ExternalPythonOperator, "_external_airflow_version_script", new_callable=mock.PropertyMock
+        ) as mock_script:
+            mock_script.return_value = "raise SystemExit('Something went wrong')"
+            caplog.set_level("WARNING")
+            caplog.clear()
+            assert op._get_airflow_version_from_target_env() is None
+            assert "Something went wrong" in caplog.text
+            assert "returned non-zero exit status" in caplog.text
+
 
 class BaseTestBranchPythonVirtualenvOperator(BaseTestPythonVirtualenvOperator):
     @pytest.fixture(autouse=True)