This is an automated email from the ASF dual-hosted git repository. ephraimanierobi pushed a commit to branch v2-6-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit ef0e84162a018526d6e3c8fd609062e9b684d47b Author: Hussein Awala <[email protected]> AuthorDate: Tue Apr 18 03:23:45 2023 +0200 rename skip_exit_code to skip_on_exit_code and allow providing multiple codes (#30692) * rename skip_exit_code to skip_on_exit_code and allow providing multiple codes * replace list type by Container (cherry picked from commit 553fb1f2ed0dfd177b11e3cc5232b964360fac8b) --- airflow/operators/python.py | 17 ++++++++++++----- tests/operators/test_python.py | 11 +++++++---- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/airflow/operators/python.py b/airflow/operators/python.py index 6f565f3492..b3744c7c85 100644 --- a/airflow/operators/python.py +++ b/airflow/operators/python.py @@ -26,6 +26,7 @@ import sys import types import warnings from abc import ABCMeta, abstractmethod +from collections.abc import Container from pathlib import Path from tempfile import TemporaryDirectory from textwrap import dedent @@ -471,7 +472,7 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator): :param expect_airflow: expect Airflow to be installed in the target environment. If true, the operator will raise warning if Airflow is not installed, and it will attempt to load Airflow macros when starting. - :param skip_exit_code: If python_callable exits with this exit code, leave the task + :param skip_on_exit_code: If python_callable exits with this exit code, leave the task in ``skipped`` state (default: None). If set to ``None``, any non-zero exit code will be treated as a failure. """ @@ -494,7 +495,7 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator): templates_dict: dict | None = None, templates_exts: list[str] | None = None, expect_airflow: bool = True, - skip_exit_code: int | None = None, + skip_on_exit_code: int | Container[int] | None = None, **kwargs, ): if ( @@ -518,7 +519,13 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator): self.python_version = python_version self.system_site_packages = system_site_packages self.pip_install_options = pip_install_options - self.skip_exit_code = skip_exit_code + self.skip_on_exit_code = ( + skip_on_exit_code + if isinstance(skip_on_exit_code, Container) + else [skip_on_exit_code] + if skip_on_exit_code + else [] + ) super().__init__( python_callable=python_callable, use_dill=use_dill, @@ -557,8 +564,8 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator): try: result = self._execute_python_callable_in_subprocess(python_path, tmp_path) except subprocess.CalledProcessError as e: - if self.skip_exit_code and e.returncode == self.skip_exit_code: - raise AirflowSkipException(f"Process exited with code {self.skip_exit_code}. Skipping.") + if e.returncode in self.skip_on_exit_code: + raise AirflowSkipException(f"Process exited with code {e.returncode}. Skipping.") else: raise return result diff --git a/tests/operators/test_python.py b/tests/operators/test_python.py index 067e07d6c3..32c3b30b64 100644 --- a/tests/operators/test_python.py +++ b/tests/operators/test_python.py @@ -938,12 +938,15 @@ class TestPythonVirtualenvOperator(BasePythonTest): "extra_kwargs, actual_exit_code, expected_state", [ (None, 99, TaskInstanceState.FAILED), - ({"skip_exit_code": 100}, 100, TaskInstanceState.SKIPPED), - ({"skip_exit_code": 100}, 101, TaskInstanceState.FAILED), - ({"skip_exit_code": None}, 0, TaskInstanceState.SUCCESS), + ({"skip_on_exit_code": 100}, 100, TaskInstanceState.SKIPPED), + ({"skip_on_exit_code": [100]}, 100, TaskInstanceState.SKIPPED), + ({"skip_on_exit_code": (100, 101)}, 100, TaskInstanceState.SKIPPED), + ({"skip_on_exit_code": 100}, 101, TaskInstanceState.FAILED), + ({"skip_on_exit_code": [100, 102]}, 101, TaskInstanceState.FAILED), + ({"skip_on_exit_code": None}, 0, TaskInstanceState.SUCCESS), ], ) - def test_skip_exit_code(self, extra_kwargs, actual_exit_code, expected_state): + def test_on_skip_exit_code(self, extra_kwargs, actual_exit_code, expected_state): def f(exit_code): if exit_code != 0: raise SystemExit(exit_code)
