Switch PythonVirtualenvOperator to venv from virtualenv package (apac…

…he#43568) The PythonVirtualenvOperator has been using virtualenv in order to support Python 2.7 and pre Python 3.5, however we do not need those any more - and we can use built-in venv instead. Fixes: apache#40420
astronomer · Nov 1, 2024 · fe0dd9b · fe0dd9b
1 parent 2c67c9f
commit fe0dd9b
Show file tree

Hide file tree

Showing 25 changed files with 372 additions and 426 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -36,7 +36,7 @@
 #                        much smaller.
 #
 # Use the same builder frontend version for everyone
-ARG AIRFLOW_EXTRAS="aiobotocore,amazon,async,celery,cncf-kubernetes,common-io,docker,elasticsearch,fab,ftp,google,google-auth,graphviz,grpc,hashicorp,http,ldap,microsoft-azure,mysql,odbc,openlineage,pandas,postgres,redis,sendgrid,sftp,slack,snowflake,ssh,statsd,uv,virtualenv"
+ARG AIRFLOW_EXTRAS="aiobotocore,amazon,async,celery,cncf-kubernetes,common-io,docker,elasticsearch,fab,ftp,google,google-auth,graphviz,grpc,hashicorp,http,ldap,microsoft-azure,mysql,odbc,openlineage,pandas,postgres,redis,sendgrid,sftp,slack,snowflake,ssh,statsd,uv"
 ARG ADDITIONAL_AIRFLOW_EXTRAS=""
 ARG ADDITIONAL_PYTHON_DEPS=""
 

diff --git a/INSTALL b/INSTALL
@@ -259,8 +259,7 @@ Those extras are available as regular core airflow extras - they install optiona
 # START CORE EXTRAS HERE
 
 aiobotocore, apache-atlas, apache-webhdfs, async, cgroups, cloudpickle, github-enterprise, google-
-auth, graphviz, kerberos, ldap, leveldb, otel, pandas, password, rabbitmq, s3fs, sentry, statsd, uv,
-virtualenv
+auth, graphviz, kerberos, ldap, leveldb, otel, pandas, password, rabbitmq, s3fs, sentry, statsd, uv
 
 # END CORE EXTRAS HERE
 

diff --git a/airflow/example_dags/example_branch_operator.py b/airflow/example_dags/example_branch_operator.py
@@ -29,143 +29,140 @@
 
 import pendulum
 
-from airflow.providers.standard.operators.python import is_venv_installed
-
-if is_venv_installed():
-    from airflow.models.dag import DAG
-    from airflow.operators.empty import EmptyOperator
-    from airflow.providers.standard.operators.python import (
-        BranchExternalPythonOperator,
-        BranchPythonOperator,
-        BranchPythonVirtualenvOperator,
-        ExternalPythonOperator,
-        PythonOperator,
-        PythonVirtualenvOperator,
+from airflow.models.dag import DAG
+from airflow.operators.empty import EmptyOperator
+from airflow.providers.standard.operators.python import (
+    BranchExternalPythonOperator,
+    BranchPythonOperator,
+    BranchPythonVirtualenvOperator,
+    ExternalPythonOperator,
+    PythonOperator,
+    PythonVirtualenvOperator,
+)
+from airflow.utils.edgemodifier import Label
+from airflow.utils.trigger_rule import TriggerRule
+
+PATH_TO_PYTHON_BINARY = sys.executable
+
+with DAG(
+    dag_id="example_branch_operator",
+    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
+    catchup=False,
+    schedule="@daily",
+    tags=["example", "example2"],
+    orientation="TB",
+) as dag:
+    run_this_first = EmptyOperator(
+        task_id="run_this_first",
     )
-    from airflow.utils.edgemodifier import Label
-    from airflow.utils.trigger_rule import TriggerRule
-
-    PATH_TO_PYTHON_BINARY = sys.executable
-
-    with DAG(
-        dag_id="example_branch_operator",
-        start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
-        catchup=False,
-        schedule="@daily",
-        tags=["example", "example2"],
-        orientation="TB",
-    ) as dag:
-        run_this_first = EmptyOperator(
-            task_id="run_this_first",
-        )
 
-        options = ["a", "b", "c", "d"]
+    options = ["a", "b", "c", "d"]
+
+    # Example branching on standard Python tasks
 
-        # Example branching on standard Python tasks
+    # [START howto_operator_branch_python]
+    branching = BranchPythonOperator(
+        task_id="branching",
+        python_callable=lambda: f"branch_{random.choice(options)}",
+    )
+    # [END howto_operator_branch_python]
+    run_this_first >> branching
+
+    join = EmptyOperator(
+        task_id="join",
+        trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
+    )
 
-        # [START howto_operator_branch_python]
-        branching = BranchPythonOperator(
-            task_id="branching",
-            python_callable=lambda: f"branch_{random.choice(options)}",
+    for option in options:
+        t = PythonOperator(
+            task_id=f"branch_{option}",
+            python_callable=lambda: print("Hello World"),
         )
-        # [END howto_operator_branch_python]
-        run_this_first >> branching
 
-        join = EmptyOperator(
-            task_id="join",
-            trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
+        empty_follow = EmptyOperator(
+            task_id="follow_" + option,
         )
 
-        for option in options:
-            t = PythonOperator(
-                task_id=f"branch_{option}",
-                python_callable=lambda: print("Hello World"),
-            )
+        # Label is optional here, but it can help identify more complex branches
+        branching >> Label(option) >> t >> empty_follow >> join
 
-            empty_follow = EmptyOperator(
-                task_id="follow_" + option,
-            )
+    # Example the same with external Python calls
 
-            # Label is optional here, but it can help identify more complex branches
-            branching >> Label(option) >> t >> empty_follow >> join
+    # [START howto_operator_branch_ext_py]
+    def branch_with_external_python(choices):
+        import random
 
-        # Example the same with external Python calls
+        return f"ext_py_{random.choice(choices)}"
 
-        # [START howto_operator_branch_ext_py]
-        def branch_with_external_python(choices):
-            import random
+    branching_ext_py = BranchExternalPythonOperator(
+        task_id="branching_ext_python",
+        python=PATH_TO_PYTHON_BINARY,
+        python_callable=branch_with_external_python,
+        op_args=[options],
+    )
+    # [END howto_operator_branch_ext_py]
+    join >> branching_ext_py
+
+    join_ext_py = EmptyOperator(
+        task_id="join_ext_python",
+        trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
+    )
 
-            return f"ext_py_{random.choice(choices)}"
+    def hello_world_with_external_python():
+        print("Hello World from external Python")
 
-        branching_ext_py = BranchExternalPythonOperator(
-            task_id="branching_ext_python",
+    for option in options:
+        t = ExternalPythonOperator(
+            task_id=f"ext_py_{option}",
             python=PATH_TO_PYTHON_BINARY,
-            python_callable=branch_with_external_python,
-            op_args=[options],
+            python_callable=hello_world_with_external_python,
         )
-        # [END howto_operator_branch_ext_py]
-        join >> branching_ext_py
 
-        join_ext_py = EmptyOperator(
-            task_id="join_ext_python",
-            trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
-        )
+        # Label is optional here, but it can help identify more complex branches
+        branching_ext_py >> Label(option) >> t >> join_ext_py
 
-        def hello_world_with_external_python():
-            print("Hello World from external Python")
+    # Example the same with Python virtual environments
 
-        for option in options:
-            t = ExternalPythonOperator(
-                task_id=f"ext_py_{option}",
-                python=PATH_TO_PYTHON_BINARY,
-                python_callable=hello_world_with_external_python,
-            )
+    # [START howto_operator_branch_virtualenv]
+    # Note: Passing a caching dir allows to keep the virtual environment over multiple runs
+    #       Run the example a second time and see that it re-uses it and is faster.
+    VENV_CACHE_PATH = Path(tempfile.gettempdir())
 
-            # Label is optional here, but it can help identify more complex branches
-            branching_ext_py >> Label(option) >> t >> join_ext_py
+    def branch_with_venv(choices):
+        import random
 
-        # Example the same with Python virtual environments
+        import numpy as np
 
-        # [START howto_operator_branch_virtualenv]
-        # Note: Passing a caching dir allows to keep the virtual environment over multiple runs
-        #       Run the example a second time and see that it re-uses it and is faster.
-        VENV_CACHE_PATH = Path(tempfile.gettempdir())
+        print(f"Some numpy stuff: {np.arange(6)}")
+        return f"venv_{random.choice(choices)}"
 
-        def branch_with_venv(choices):
-            import random
+    branching_venv = BranchPythonVirtualenvOperator(
+        task_id="branching_venv",
+        requirements=["numpy~=1.26.0"],
+        venv_cache_path=VENV_CACHE_PATH,
+        python_callable=branch_with_venv,
+        op_args=[options],
+    )
+    # [END howto_operator_branch_virtualenv]
+    join_ext_py >> branching_venv
+
+    join_venv = EmptyOperator(
+        task_id="join_venv",
+        trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
+    )
 
-            import numpy as np
+    def hello_world_with_venv():
+        import numpy as np
 
-            print(f"Some numpy stuff: {np.arange(6)}")
-            return f"venv_{random.choice(choices)}"
+        print(f"Hello World with some numpy stuff: {np.arange(6)}")
 
-        branching_venv = BranchPythonVirtualenvOperator(
-            task_id="branching_venv",
+    for option in options:
+        t = PythonVirtualenvOperator(
+            task_id=f"venv_{option}",
             requirements=["numpy~=1.26.0"],
             venv_cache_path=VENV_CACHE_PATH,
-            python_callable=branch_with_venv,
-            op_args=[options],
-        )
-        # [END howto_operator_branch_virtualenv]
-        join_ext_py >> branching_venv
-
-        join_venv = EmptyOperator(
-            task_id="join_venv",
-            trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
+            python_callable=hello_world_with_venv,
         )
 
-        def hello_world_with_venv():
-            import numpy as np
-
-            print(f"Hello World with some numpy stuff: {np.arange(6)}")
-
-        for option in options:
-            t = PythonVirtualenvOperator(
-                task_id=f"venv_{option}",
-                requirements=["numpy~=1.26.0"],
-                venv_cache_path=VENV_CACHE_PATH,
-                python_callable=hello_world_with_venv,
-            )
-
-            # Label is optional here, but it can help identify more complex branches
-            branching_venv >> Label(option) >> t >> join_venv
+        # Label is optional here, but it can help identify more complex branches
+        branching_venv >> Label(option) >> t >> join_venv