This is an automated email from the ASF dual-hosted git repository. ash pushed a commit to branch v2-0-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit b583736a1ee909e5c2d5255ae04870e3ad434c06 Author: penggongkui <[email protected]> AuthorDate: Thu Mar 18 19:24:11 2021 +0800 [AIRFLOW-6076] fix dag.cli() KeyError (#13647) (cherry picked from commit b24a1babd4271d74ba13b1dc0a9cf8da001b3f77) --- airflow/cli/cli_parser.py | 33 +++++++++++++++++++++++++++++---- tests/cli/test_cli_parser.py | 26 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index b198234..fae516e 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -24,7 +24,7 @@ import os import textwrap from argparse import Action, ArgumentError, RawTextHelpFormatter from functools import lru_cache -from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Union +from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Union from airflow import settings from airflow.cli.commands.legacy_commands import check_legacy_command @@ -1511,7 +1511,31 @@ airflow_commands: List[CLICommand] = [ ), ] ALL_COMMANDS_DICT: Dict[str, CLICommand] = {sp.name: sp for sp in airflow_commands} -DAG_CLI_COMMANDS: Set[str] = {'list_tasks', 'backfill', 'test', 'run', 'pause', 'unpause', 'list_dag_runs'} + + +def _remove_dag_id_opt(command: ActionCommand): + cmd = command._asdict() + cmd['args'] = (arg for arg in command.args if arg is not ARG_DAG_ID) + return ActionCommand(**cmd) + + +dag_cli_commands: List[CLICommand] = [ + GroupCommand( + name='dags', + help='Manage DAGs', + subcommands=[ + _remove_dag_id_opt(sp) + for sp in DAGS_COMMANDS + if sp.name in ['backfill', 'list-runs', 'pause', 'unpause'] + ], + ), + GroupCommand( + name='tasks', + help='Manage tasks', + subcommands=[_remove_dag_id_opt(sp) for sp in TASKS_COMMANDS if sp.name in ['list', 'test', 'run']], + ), +] +DAG_CLI_DICT: Dict[str, CLICommand] = {sp.name: sp for sp in dag_cli_commands} class AirflowHelpFormatter(argparse.HelpFormatter): @@ -1563,10 +1587,11 @@ def get_parser(dag_parser: bool = False) -> argparse.ArgumentParser: subparsers = parser.add_subparsers(dest='subcommand', metavar="GROUP_OR_COMMAND") subparsers.required = True - subparser_list = DAG_CLI_COMMANDS if dag_parser else ALL_COMMANDS_DICT.keys() + command_dict = DAG_CLI_DICT if dag_parser else ALL_COMMANDS_DICT + subparser_list = command_dict.keys() sub_name: str for sub_name in sorted(subparser_list): - sub: CLICommand = ALL_COMMANDS_DICT[sub_name] + sub: CLICommand = command_dict[sub_name] _add_command(subparsers, sub) return parser diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py index 1c2e2aa..0ba5dff 100644 --- a/tests/cli/test_cli_parser.py +++ b/tests/cli/test_cli_parser.py @@ -144,6 +144,16 @@ class TestCli(TestCase): assert "Commands" in stdout assert "Groups" in stdout + def test_dag_parser_commands_and_comamnd_group_sections(self): + parser = cli_parser.get_parser(dag_parser=True) + + with contextlib.redirect_stdout(io.StringIO()) as stdout: + with self.assertRaises(SystemExit): + parser.parse_args(['--help']) + stdout = stdout.getvalue() + self.assertIn("Commands", stdout) + self.assertIn("Groups", stdout) + def test_should_display_help(self): parser = cli_parser.get_parser() @@ -160,6 +170,22 @@ class TestCli(TestCase): with pytest.raises(SystemExit): parser.parse_args([*cmd_args, '--help']) + def test_dag_cli_should_display_help(self): + parser = cli_parser.get_parser(dag_parser=True) + + all_command_as_args = [ + command_as_args + for top_command in cli_parser.dag_cli_commands + for command_as_args in ( + [[top_command.name]] + if isinstance(top_command, cli_parser.ActionCommand) + else [[top_command.name, nested_command.name] for nested_command in top_command.subcommands] + ) + ] + for cmd_args in all_command_as_args: + with self.assertRaises(SystemExit): + parser.parse_args([*cmd_args, '--help']) + def test_positive_int(self): assert 1 == cli_parser.positive_int('1')
