Skip to content

Commit fb8b590

Browse files
ashbamoghrajesh
andauthored
Allow virtualenv code to access connections/variables and send logs (apache#57213)
This change is quite simple, as it simply calls the function added in apache#57212 if it's available. Up until the point that it is released this code will not hit the `getattr` will return None and be a no-op. I didn't use the walrus operator here as _theoretically_ someone could be using this to run in a Python that doesn't support it. Fixes apache#51422, fixes apache#54706 Co-authored-by: Amogh Desai <amoghrajesh1999@gmail.com>
1 parent ff4e471 commit fb8b590

File tree

1 file changed

+11
-3
lines changed

1 file changed

+11
-3
lines changed

providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,16 @@ if sys.version_info >= (3,6):
3838
{# Airflow is not available in this environment, therefore we won't
3939
# be able to integrate any plugin macros. #}
4040
pass
41+
42+
try:
43+
from airflow.sdk.execution_time import task_runner
44+
except ModuleNotFoundError:
45+
pass
46+
else:
47+
{#- We are in an Airflow 3.x env, try and set up supervisor comms so virtual env can still access tasks etc! #}
48+
reinit_supervisor_comms = getattr(task_runner, "reinit_supervisor_comms", None)
49+
if reinit_supervisor_comms:
50+
reinit_supervisor_comms()
4151
{% endif %}
4252

4353
# Script
@@ -49,12 +59,10 @@ if sys.version_info >= (3,6):
4959
import types
5060

5161
{{ modified_dag_module_name }} = types.ModuleType("{{ modified_dag_module_name }}")
52-
5362
{{ modified_dag_module_name }}.{{ python_callable }} = {{ python_callable }}
54-
5563
sys.modules["{{modified_dag_module_name}}"] = {{modified_dag_module_name}}
5664

57-
{% endif%}
65+
{%- endif -%}
5866

5967
{% if op_args or op_kwargs %}
6068
with open(sys.argv[1], "rb") as file:

0 commit comments

Comments
 (0)