From c8ab8588a45027c6e3b036a4d376c3cacd8aebf2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Mon, 19 Jul 2021 02:08:40 +0200 Subject: [PATCH 1/3] Improve executor validation in CLI --- airflow/cli/cli_parser.py | 13 ++++++++----- tests/cli/test_cli_parser.py | 22 ++++++++++++++++++++++ 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index ecf3b18d44bb3..b6528384d1fa0 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -57,11 +57,14 @@ class DefaultHelpParser(argparse.ArgumentParser): def _check_value(self, action, value): """Override _check_value and check conditionally added command""" - executor = conf.get('core', 'EXECUTOR') - if value == 'celery' and executor not in (CELERY_EXECUTOR, CELERY_KUBERNETES_EXECUTOR): - message = f'celery subcommand works only with CeleryExecutor, your current executor: {executor}' - raise ArgumentError(action, message) - if value == 'kubernetes': + if action.dest == 'subcommand' and value == 'celery': + executor = conf.get('core', 'EXECUTOR') + if executor not in (CELERY_EXECUTOR, CELERY_KUBERNETES_EXECUTOR): + message = ( + f'celery subcommand works only with CeleryExecutor, your current executor: {executor}' + ) + raise ArgumentError(action, message) + if action.dest == 'subcommand' and value == 'kubernetes': try: import kubernetes.client # noqa: F401 except ImportError: diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index e6cb9bc82dae0..7cece6d535e5d 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -25,10 +25,13 @@ from unittest import TestCase import pytest +from parameterized import parameterized from airflow.cli import cli_parser # Can not be `--snake_case` or contain uppercase letter +from tests.test_utils.config import conf_vars + ILLEGAL_LONG_OPTION_PATTERN = re.compile("^--[a-z]+_[a-z]+|^--.*[A-Z].*") # Only can be `-[a-z]` or `-[A-Z]` LEGAL_SHORT_OPTION_PATTERN = re.compile("^-[a-zA-z]$") @@ -193,3 +196,22 @@ def test_positive_int(self): with pytest.raises(argparse.ArgumentTypeError): cli_parser.positive_int(allow_zero=False)('0') cli_parser.positive_int(allow_zero=True)('-1') + + @parameterized.expand(["CeleryExecutor", "CeleryKubernetesExecutor"]) + def test_dag_parser_celery_command_required_celery_executor(self, executor): + with conf_vars({('core', 'executor'): executor}), contextlib.redirect_stderr(io.StringIO()) as stderr: + parser = cli_parser.get_parser() + with self.assertRaises(SystemExit): + parser.parse_args(['celery']) + stderr = stderr.getvalue() + assert ( + "airflow celery command error: the following arguments are required: COMMAND, see help above." + ) in stderr + + def test_dag_parser_config_command_dont_required_celery_executor(self): + with conf_vars({('core', 'executor'): "CeleryExecutor"}), contextlib.redirect_stderr( + io.StringIO() + ) as stdout: + parser = cli_parser.get_parser() + parser.parse_args(['config', 'get-value', 'celery', 'broker-url']) + assert stdout is not None From 374016f245c91368b9d6a06c3e0dca576a635a99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Mon, 19 Jul 2021 02:12:06 +0200 Subject: [PATCH 2/3] fixup! Improve executor validation in CLI --- tests/cli/test_cli_parser.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index 7cece6d535e5d..b643a9e27e8b7 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -28,10 +28,9 @@ from parameterized import parameterized from airflow.cli import cli_parser - -# Can not be `--snake_case` or contain uppercase letter from tests.test_utils.config import conf_vars +# Can not be `--snake_case` or contain uppercase letter ILLEGAL_LONG_OPTION_PATTERN = re.compile("^--[a-z]+_[a-z]+|^--.*[A-Z].*") # Only can be `-[a-z]` or `-[A-Z]` LEGAL_SHORT_OPTION_PATTERN = re.compile("^-[a-zA-z]$") From d8dea8a4288e874167689c078d295117b35c06eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kamil=20Bregu=C5=82a?= Date: Mon, 19 Jul 2021 02:26:56 +0200 Subject: [PATCH 3/3] fixup! fixup! Improve executor validation in CLI --- tests/cli/test_cli_parser.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index b643a9e27e8b7..55ec4927bc0e7 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -196,8 +196,21 @@ def test_positive_int(self): cli_parser.positive_int(allow_zero=False)('0') cli_parser.positive_int(allow_zero=True)('-1') + def test_dag_parser_celery_command_require_celery_executor(self): + with conf_vars({('core', 'executor'): 'SequentialExecutor'}), contextlib.redirect_stderr( + io.StringIO() + ) as stderr: + parser = cli_parser.get_parser() + with self.assertRaises(SystemExit): + parser.parse_args(['celery']) + stderr = stderr.getvalue() + assert ( + "airflow command error: argument GROUP_OR_COMMAND: celery subcommand " + "works only with CeleryExecutor, your current executor: SequentialExecutor, see help above." + ) in stderr + @parameterized.expand(["CeleryExecutor", "CeleryKubernetesExecutor"]) - def test_dag_parser_celery_command_required_celery_executor(self, executor): + def test_dag_parser_celery_command_accept_celery_executor(self, executor): with conf_vars({('core', 'executor'): executor}), contextlib.redirect_stderr(io.StringIO()) as stderr: parser = cli_parser.get_parser() with self.assertRaises(SystemExit):