diff --git a/airflow-core/tests/unit/always/test_example_dags.py b/airflow-core/tests/unit/always/test_example_dags.py index b45845404ae53..40763326c8698 100644 --- a/airflow-core/tests/unit/always/test_example_dags.py +++ b/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(): diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 4708365e846e2..6bc89c7b4073e 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/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 diff --git a/providers/standard/docs/operators/python.rst b/providers/standard/docs/operators/python.rst index 30f7151e55bef..1e3c7bd3fc445 100644 --- a/providers/standard/docs/operators/python.rst +++ b/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. diff --git a/providers/standard/src/airflow/providers/standard/operators/python.py b/providers/standard/src/airflow/providers/standard/operators/python.py index 8698ddc403b60..9121cbfaead99 100644 --- a/providers/standard/src/airflow/providers/standard/operators/python.py +++ b/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, ) diff --git a/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 b/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 index 4a6c281463337..32b88574e47e7 100644 --- a/providers/standard/src/airflow/providers/standard/utils/python_virtualenv_script.jinja2 +++ b/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: diff --git a/providers/standard/tests/unit/standard/operators/test_python.py b/providers/standard/tests/unit/standard/operators/test_python.py index d8228b1032f30..7f01e7ee07efb 100644 --- a/providers/standard/tests/unit/standard/operators/test_python.py +++ b/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():