Restore Airflow context in virtualenv operators

airflow-core/tests/unit/always/test_example_dags.py

@@ -149,12 +149,6 @@ def example_not_excluded_dags(xfail_db_exception: bool = False):
                     pytest.mark.skip(reason=f"Not supported for Python {CURRENT_PYTHON_VERSION}")
                 )
 
-            # TODO: remove when context serialization is implemented in AIP-72
-            if "/example_python_context_" in candidate:
-                param_marks.append(
-                    pytest.mark.skip(reason="Temporary excluded until AIP-72 context serialization is done.")
-                )
-
             for optional, dependencies in OPTIONAL_PROVIDERS_DEPENDENCIES.items():
                 if re.match(optional, candidate):
                     for distribution_name, specifier in dependencies.items():

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -567,8 +567,6 @@ def test_serialization(self):
             # AirflowProviderDeprecationWarning etc show up in import_errors, and being aware of all of those is
             # not relevant to this test; we only care about actual errors
             if "airflow.exceptions.AirflowProviderDeprecationWarning" not in error
-            # TODO: TaskSDK
-            if "`use_airflow_context=True` is not yet implemented" not in error
             # This "looks" like a problem, but is just a quirk of the parse-all-dags-in-one-process we do
             # in this test
             if "AirflowDagDuplicatedIdException: Ignoring DAG example_sagemaker" not in error

providers/standard/docs/operators/python.rst

@@ -193,6 +193,21 @@ setting ``system_site_packages`` to ``True`` or you won't have access to most co
 If you want the context related to datetime objects like ``data_interval_start``, you can add ``pendulum`` and
 ``lazy_object_proxy``.
 
+Accessing the current context
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+On Airflow 3, set ``use_airflow_context=True`` to make the task context available through
+:func:`airflow.sdk.get_current_context` inside the virtual environment. Airflow must be installed in that
+environment, either through ``system_site_packages=True`` or as one of the environment's requirements.
+
+.. code-block:: python
+
+    @task.virtualenv(system_site_packages=True, use_airflow_context=True)
+    def print_run_id():
+        from airflow.sdk import get_current_context
+
+        print(get_current_context()["run_id"])
+
 .. important::
 
     When Airflow or provider packages are required, you must specify the Airflow :ref:`apache-airflow:installation:constraints`
@@ -319,6 +334,21 @@ Otherwise you won't have access to the most context variables of Airflow in ``op
 If you want the context related to datetime objects like ``data_interval_start`` you can add ``pendulum`` and
 ``lazy_object_proxy`` to your virtual environment.
 
+Accessing the current context
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+On Airflow 3, set ``use_airflow_context=True`` to make the task context available through
+:func:`airflow.sdk.get_current_context`. The external environment must contain the same Airflow version as the
+main environment, and ``expect_airflow`` must remain enabled.
+
+.. code-block:: python
+
+    @task.external_python(python=PATH_TO_PYTHON_BINARY, use_airflow_context=True)
+    def print_run_id():
+        from airflow.sdk import get_current_context
+
+        print(get_current_context()["run_id"])
+
 .. important::
     The Python function body defined to be executed is cut out of the Dag into a temporary file w/o surrounding code.
     As in the examples you need to add all imports again and you can not rely on variables from the global Python context.

providers/standard/src/airflow/providers/standard/operators/python.py

@@ -499,6 +499,7 @@ def __init__(
         skip_on_exit_code: int | Container[int] | None = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
+        use_airflow_context: bool = False,
         **kwargs,
     ):
         if (
@@ -539,6 +540,9 @@ def __init__(
         )
         self.env_vars = env_vars
         self.inherit_env = inherit_env
+        if use_airflow_context and not AIRFLOW_V_3_0_PLUS:
+            raise ValueError("use_airflow_context requires Airflow 3.0 or later.")
+        self.use_airflow_context = use_airflow_context
 
     @abstractmethod
     def _iter_serializable_context_keys(self):
@@ -548,6 +552,7 @@ def execute(self, context: Context) -> Any:
         serializable_keys = set(self._iter_serializable_context_keys())
         new = {k: v for k, v in context.items() if k in serializable_keys}
         serializable_context = cast("Context", new)
+        self._airflow_context = serializable_context
         # Store bundle_path for subprocess execution
         self._bundle_path = self._get_bundle_path_from_context(context)
         return super().execute(context=serializable_context)
@@ -639,10 +644,15 @@ def _execute_python_callable_in_subprocess(self, python_path: Path):
             string_args_path = tmp_dir / "string_args.txt"
             script_path = tmp_dir / "script.py"
             termination_log_path = tmp_dir / "termination.log"
-            airflow_context_path = tmp_dir / "airflow_context.json"
+            airflow_context_path = tmp_dir / "airflow_context.bin"
 
             self._write_args(input_path)
             self._write_string_args(string_args_path)
+            if self.use_airflow_context:
+                from airflow.serialization.serialized_objects import BaseSerialization
+
+                serialized_context = BaseSerialization.serialize(self._airflow_context)
+                airflow_context_path.write_bytes(self.pickling_library.dumps(serialized_context))
 
             jinja_context = {
                 "op_args": self.op_args,
@@ -651,6 +661,7 @@ def _execute_python_callable_in_subprocess(self, python_path: Path):
                 "pickling_library": self.serializer,
                 "python_callable": self.python_callable.__name__,
                 "python_callable_source": self.get_python_source(),
+                "use_airflow_context": self.use_airflow_context,
                 **self._get_additional_jinja_context(),
             }
 
@@ -827,6 +838,9 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
         environment. If set to ``True``, the virtual environment will inherit the environment variables
         of the parent process (``os.environ``). If set to ``False``, the virtual environment will be
         executed with a clean environment.
+    :param use_airflow_context: Whether to make the current Airflow context available to
+        ``get_current_context()`` in the virtual environment. Airflow must be installed in the
+        virtual environment.
     """
 
     template_fields: Sequence[str] = tuple(
@@ -857,6 +871,7 @@ def __init__(
         venv_cache_path: None | os.PathLike[str] = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
+        use_airflow_context: bool = False,
         **kwargs,
     ):
         if (
@@ -873,6 +888,11 @@ def __init__(
             raise AirflowException(
                 "Passing non-string types (e.g. int or float) as python_version not supported"
             )
+        if use_airflow_context and not (expect_airflow or system_site_packages):
+            raise ValueError(
+                "use_airflow_context requires Airflow in the virtual environment; set "
+                "expect_airflow or system_site_packages to True."
+            )
         if not requirements:
             self.requirements: list[str] = []
         elif isinstance(requirements, str):
@@ -907,6 +927,7 @@ def __init__(
             skip_on_exit_code=skip_on_exit_code,
             env_vars=env_vars,
             inherit_env=inherit_env,
+            use_airflow_context=use_airflow_context,
             **kwargs,
         )
 
@@ -1220,6 +1241,9 @@ class ExternalPythonOperator(_BasePythonVirtualenvOperator):
         environment. If set to ``True``, the virtual environment will inherit the environment variables
         of the parent process (``os.environ``). If set to ``False``, the virtual environment will be
         executed with a clean environment.
+    :param use_airflow_context: Whether to make the current Airflow context available to
+        ``get_current_context()`` in the external Python environment. Airflow must be installed in the
+        external environment.
     """
 
     template_fields: Sequence[str] = tuple({"python"}.union(PythonOperator.template_fields))
@@ -1240,10 +1264,13 @@ def __init__(
         skip_on_exit_code: int | Container[int] | None = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
+        use_airflow_context: bool = False,
         **kwargs,
     ):
         if not python:
             raise ValueError("Python Path must be defined in ExternalPythonOperator")
+        if use_airflow_context and not expect_airflow:
+            raise ValueError("use_airflow_context requires expect_airflow to be True.")
         self.python = python
         self.expect_pendulum = expect_pendulum
         super().__init__(
@@ -1258,6 +1285,7 @@ def __init__(
             skip_on_exit_code=skip_on_exit_code,
             env_vars=env_vars,
             inherit_env=inherit_env,
+            use_airflow_context=use_airflow_context,
             **kwargs,
         )
 

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

@@ -128,12 +128,27 @@ with open(sys.argv[3], "r") as file:
     virtualenv_string_args = list(map(lambda x: x.strip(), list(file)))
 {% endif %}
 
+{% if use_airflow_context | default(false) -%}
+from airflow.sdk.execution_time.context import set_current_context
+from airflow.serialization.serialized_objects import BaseSerialization
+
+with open(sys.argv[5], "rb") as file:
+    airflow_context = BaseSerialization.deserialize({{ pickling_library }}.load(file))
+with set_current_context(airflow_context):
+    try:
+        res = {{ python_callable }}(*arg_dict["args"], **arg_dict["kwargs"])
+    except Exception as e:
+        with open(sys.argv[4], "w") as file:
+            file.write(str(e))
+        raise
+{% else %}
 try:
     res = {{ python_callable }}(*arg_dict["args"], **arg_dict["kwargs"])
 except Exception as e:
     with open(sys.argv[4], "w") as file:
         file.write(str(e))
     raise
+{% endif %}
 
 # Write output
 with open(sys.argv[2], "wb") as file:

providers/standard/tests/unit/standard/operators/test_python.py

@@ -211,6 +211,11 @@ def _pull_xcom(ran):
         """
         return TaskInstance.xcom_pull(ran)
 
+    @staticmethod
+    def _run_id(ran):
+        """Return the run ID for whatever ``run_as_task(return_ti=True)`` returned."""
+        return ran.run_id
+
     def render_templates(self, fn, **kwargs):
         """Create TaskInstance and render templates without actual run."""
         return self.create_ti(fn, **kwargs).render_templates()
@@ -1367,6 +1372,10 @@ def run_as_operator(self, fn, **kwargs):
     def _pull_xcom(self, ran):  # ``ran`` is the TaskRunResult from run_as_task(return_ti=True)
         return pushed_xcom(self._last_xcoms, ran.ti)
 
+    @staticmethod
+    def _run_id(ran):
+        return ran.ti.run_id
+
 
 def _dbfree_venv_supported() -> bool:
     """Whether venv operators can run DB-free here.
@@ -1424,6 +1433,61 @@ def default_kwargs(*, python_version=DEFAULT_PYTHON_VERSION, **kwargs):
                 kwargs["venv_cache_path"] = venv_cache_path
         return kwargs
 
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="The Task SDK is only available in Airflow 3")
+    def test_get_current_context(self):
+        def f():
+            from airflow.sdk import get_current_context
+
+            context = get_current_context()
+            return {"run_id": context["run_id"], "task_id": context["task"].task_id}
+
+        result = self.run_as_task(
+            f,
+            env_vars={"PYTHONPATH": os.pathsep.join(sys.path)},
+            return_ti=True,
+            system_site_packages=True,
+            use_airflow_context=True,
+        )
+
+        assert self._pull_xcom(result) == {"run_id": self._run_id(result), "task_id": self.task_id}
+
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="The Task SDK is only available in Airflow 3")
+    def test_get_current_context_disabled_by_default(self):
+        def f():
+            from airflow.sdk import get_current_context
+
+            get_current_context()
+
+        with pytest.raises(AirflowException, match="Current context was requested but no context was found"):
+            self.run_as_task(
+                f,
+                env_vars={"PYTHONPATH": os.pathsep.join(sys.path)},
+                system_site_packages=True,
+            )
+
+    def test_use_airflow_context_requires_airflow_in_virtualenv(self):
+        def f(): ...
+
+        with pytest.raises(ValueError, match="requires Airflow in the virtual environment"):
+            PythonVirtualenvOperator(
+                task_id=self.task_id,
+                python_callable=f,
+                expect_airflow=False,
+                system_site_packages=False,
+                use_airflow_context=True,
+            )
+
+    @mock.patch("airflow.providers.standard.operators.python.AIRFLOW_V_3_0_PLUS", False)
+    def test_use_airflow_context_requires_airflow_3(self):
+        def f(): ...
+
+        with pytest.raises(ValueError, match="requires Airflow 3.0 or later"):
+            PythonVirtualenvOperator(
+                task_id=self.task_id,
+                python_callable=f,
+                use_airflow_context=True,
+            )
+
     @CLOUDPICKLE_MARKER
     def test_add_cloudpickle(self):
         def f():
@@ -2107,6 +2171,30 @@ def default_kwargs(*, python_version=DEFAULT_PYTHON_VERSION, **kwargs):
         kwargs["python"] = sys.executable
         return kwargs
 
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="The Task SDK is only available in Airflow 3")
+    def test_get_current_context(self):
+        def f():
+            from airflow.sdk import get_current_context
+
+            context = get_current_context()
+            return {"run_id": context["run_id"], "task_id": context["task"].task_id}
+
+        result = self.run_as_task(f, return_ti=True, use_airflow_context=True)
+
+        assert self._pull_xcom(result) == {"run_id": self._run_id(result), "task_id": self.task_id}
+
+    def test_use_airflow_context_requires_expect_airflow(self):
+        def f(): ...
+
+        with pytest.raises(ValueError, match="requires expect_airflow to be True"):
+            ExternalPythonOperator(
+                task_id=self.task_id,
+                python=sys.executable,
+                python_callable=f,
+                expect_airflow=False,
+                use_airflow_context=True,
+            )
+
     @mock.patch("pickle.loads")
     def test_except_value_error(self, loads_mock):
         def f():