[airflow] 03/42: [AIRFLOW-6076] fix dag.cli() KeyError (#13647)

[as...@apache.org]

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 <49...@users.noreply.github.com>
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')