Refactor list rendering in commands (#12704)

This commit unifies the mechanism of rendering output of tabular
data. This gives users a possibility to eiter display a tabular
representation of data or render it as valid json or yaml payload.

Closes: #12699

Co-authored-by: Kaxil Naik <[email protected]>

Author: turbaszek

UPDATING.md

@@ -52,6 +52,40 @@ assists users migrating to a new version.
 
 ## Master
 
+### Changes to output argument in commands
+
+From Airflow 2.0, We are replacing [tabulate](https://pypi.org/project/tabulate/) with [rich](https://github.com/willmcgugan/rich) to render commands output. Due to this change, the `--output` argument
+will no longer accept formats of tabulate tables. Instead, it now accepts:
+
+- `table` - will render the output in predefined table
+- `json` - will render the output as a json
+- `yaml` - will render the output as yaml
+
+By doing this we increased consistency and gave users possibility to manipulate the
+output programmatically (when using json or yaml).
+
+Affected commands:
+
+- `airflow dags list`
+- `airflow dags report`
+- `airflow dags list-runs`
+- `airflow dags list-jobs`
+- `airflow connections list`
+- `airflow connections get`
+- `airflow pools list`
+- `airflow pools get`
+- `airflow pools set`
+- `airflow pools delete`
+- `airflow pools import`
+- `airflow pools export`
+- `airflow role list`
+- `airflow providers list`
+- `airflow providers get`
+- `airflow providers hooks`
+- `airflow tasks states-for-dag-run`
+- `airflow users list`
+- `airflow variables list`
+
 ### Azure Wasb Hook does not work together with Snowflake hook
 
 The WasbHook in Apache Airflow use a legacy version of Azure library. While the conflict is not

airflow/cli/cli_parser.py

@@ -26,8 +26,6 @@
 from functools import lru_cache
 from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Union
 
-from tabulate import tabulate_formats
-
 from airflow import settings
 from airflow.cli.commands.legacy_commands import check_legacy_command
 from airflow.configuration import conf
@@ -176,13 +174,10 @@ def positive_int(value):
 )
 ARG_OUTPUT = Arg(
     ("--output",),
-    help=(
-        "Output table format. The specified value is passed to "
-        "the tabulate module (https://pypi.org/project/tabulate/). "
-    ),
-    metavar="FORMAT",
-    choices=tabulate_formats,
-    default="plain",
+    help=("Output format. Allowed values: json, yaml, table (default: table)"),
+    metavar="(table, json, yaml)",
+    choices=("table", "json", "yaml"),
+    default="table",
 )
 ARG_COLOR = Arg(
     ('--color',),
@@ -1033,7 +1028,7 @@ class GroupCommand(NamedTuple):
         name='list',
         help='List variables',
         func=lazy_load_command('airflow.cli.commands.variable_command.variables_list'),
-        args=(),
+        args=(ARG_OUTPUT,),
     ),
     ActionCommand(
         name='get',
@@ -1110,7 +1105,7 @@ class GroupCommand(NamedTuple):
         name='get',
         help='Get a connection',
         func=lazy_load_command('airflow.cli.commands.connection_command.connections_get'),
-        args=(ARG_CONN_ID, ARG_COLOR),
+        args=(ARG_CONN_ID, ARG_COLOR, ARG_OUTPUT),
     ),
     ActionCommand(
         name='list',

airflow/cli/commands/cheat_sheet_command.py

@@ -21,9 +21,11 @@
 
 from airflow.cli.cli_parser import ActionCommand, GroupCommand, airflow_commands
 from airflow.cli.simple_table import SimpleTable
+from airflow.utils.cli import suppress_logs_and_warning
 from airflow.utils.helpers import partition
 
 
+@suppress_logs_and_warning()
 def cheat_sheet(args):
     """Display cheat-sheet."""
     display_commands_index()

airflow/cli/commands/connection_command.py

@@ -19,78 +19,54 @@
 import json
 import os
 import sys
-from typing import List
+from typing import Any, Dict, List
 from urllib.parse import urlparse, urlunparse
 
-import pygments
 import yaml
-from pygments.lexers.data import YamlLexer
 from sqlalchemy.orm import exc
-from tabulate import tabulate
 
+from airflow.cli.simple_table import AirflowConsole
 from airflow.exceptions import AirflowNotFoundException
 from airflow.hooks.base_hook import BaseHook
 from airflow.models import Connection
 from airflow.utils import cli as cli_utils
-from airflow.utils.cli import should_use_colors
-from airflow.utils.code_utils import get_terminal_formatter
+from airflow.utils.cli import suppress_logs_and_warning
 from airflow.utils.session import create_session
 
 
-def _tabulate_connection(conns: List[Connection], tablefmt: str):
-    tabulate_data = [
-        {
-            'Conn Id': conn.conn_id,
-            'Conn Type': conn.conn_type,
-            'Description': conn.description,
-            'Host': conn.host,
-            'Port': conn.port,
-            'Is Encrypted': conn.is_encrypted,
-            'Is Extra Encrypted': conn.is_encrypted,
-            'Extra': conn.extra,
-        }
-        for conn in conns
-    ]
-
-    msg = tabulate(tabulate_data, tablefmt=tablefmt, headers='keys')
-    return msg
-
-
-def _yamulate_connection(conn: Connection):
-    yaml_data = {
-        'Id': conn.id,
-        'Conn Id': conn.conn_id,
-        'Conn Type': conn.conn_type,
-        'Description': conn.description,
-        'Host': conn.host,
-        'Schema': conn.schema,
-        'Login': conn.login,
-        'Password': conn.password,
-        'Port': conn.port,
-        'Is Encrypted': conn.is_encrypted,
-        'Is Extra Encrypted': conn.is_encrypted,
-        'Extra': conn.extra_dejson,
-        'URI': conn.get_uri(),
+def _connection_mapper(conn: Connection) -> Dict[str, Any]:
+    return {
+        'id': conn.id,
+        'conn_id': conn.conn_id,
+        'conn_type': conn.conn_type,
+        'description': conn.description,
+        'host': conn.host,
+        'schema': conn.schema,
+        'login': conn.login,
+        'password': conn.password,
+        'port': conn.port,
+        'is_encrypted': conn.is_encrypted,
+        'is_extra_encrypted': conn.is_encrypted,
+        'extra_dejson': conn.extra_dejson,
+        'get_uri': conn.get_uri(),
     }
-    return yaml.safe_dump(yaml_data, sort_keys=False)
 
 
+@suppress_logs_and_warning()
 def connections_get(args):
     """Get a connection."""
     try:
         conn = BaseHook.get_connection(args.conn_id)
     except AirflowNotFoundException:
         raise SystemExit("Connection not found.")
-
-    yaml_content = _yamulate_connection(conn)
-    if should_use_colors(args):
-        yaml_content = pygments.highlight(
-            code=yaml_content, formatter=get_terminal_formatter(), lexer=YamlLexer()
-        )
-
-    print(yaml_content)
+    AirflowConsole().print_as(
+        data=[conn],
+        output=args.output,
+        mapper=_connection_mapper,
+    )
 
 
+@suppress_logs_and_warning()
 def connections_list(args):
     """Lists all connections at the command line"""
     with create_session() as session:
@@ -99,9 +75,11 @@ def connections_list(args):
             query = query.filter(Connection.conn_id == args.conn_id)
         conns = query.all()
 
-        tablefmt = args.output
-        msg = _tabulate_connection(conns, tablefmt)
-        print(msg)
+        AirflowConsole().print_as(
+            data=conns,
+            output=args.output,
+            mapper=_connection_mapper,
+        )
 
 
 def _format_connections(conns: List[Connection], fmt: str) -> str:

airflow/cli/commands/dag_command.py

@@ -16,60 +16,38 @@
 # under the License.
 
 """Dag sub-commands"""
+import ast
 import errno
 import json
 import logging
 import signal
 import subprocess
 import sys
-from typing import List
 
 from graphviz.dot import Dot
-from tabulate import tabulate
 
 from airflow import settings
 from airflow.api.client import get_current_api_client
+from airflow.cli.simple_table import AirflowConsole
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, BackfillUnfinished
 from airflow.executors.debug_executor import DebugExecutor
 from airflow.jobs.base_job import BaseJob
 from airflow.models import DagBag, DagModel, DagRun, TaskInstance
 from airflow.models.dag import DAG
 from airflow.utils import cli as cli_utils
-from airflow.utils.cli import get_dag, get_dag_by_file_location, process_subdir, sigint_handler
+from airflow.utils.cli import (
+    get_dag,
+    get_dag_by_file_location,
+    process_subdir,
+    sigint_handler,
+    suppress_logs_and_warning,
+)
 from airflow.utils.dot_renderer import render_dag
 from airflow.utils.session import create_session, provide_session
 from airflow.utils.state import State
 
 
-def _tabulate_dag_runs(dag_runs: List[DagRun], tablefmt: str = "fancy_grid") -> str:
-    tabulate_data = (
-        {
-            'ID': dag_run.id,
-            'Run ID': dag_run.run_id,
-            'State': dag_run.state,
-            'DAG ID': dag_run.dag_id,
-            'Execution date': dag_run.execution_date.isoformat(),
-            'Start date': dag_run.start_date.isoformat() if dag_run.start_date else '',
-            'End date': dag_run.end_date.isoformat() if dag_run.end_date else '',
-        }
-        for dag_run in dag_runs
-    )
-    return tabulate(tabular_data=tabulate_data, tablefmt=tablefmt)
-
-
-def _tabulate_dags(dags: List[DAG], tablefmt: str = "fancy_grid") -> str:
-    tabulate_data = (
-        {
-            'DAG ID': dag.dag_id,
-            'Filepath': dag.filepath,
-            'Owner': dag.owner,
-        }
-        for dag in sorted(dags, key=lambda d: d.dag_id)
-    )
-    return tabulate(tabular_data=tabulate_data, tablefmt=tablefmt, headers='keys')
-
-
 @cli_utils.action_logging
 def dag_backfill(args, dag=None):
     """Creates backfill job or dry run for a DAG"""
@@ -293,21 +271,41 @@ def dag_next_execution(args):
 
 
 @cli_utils.action_logging
+@suppress_logs_and_warning()
 def dag_list_dags(args):
     """Displays dags with or without stats at the command line"""
     dagbag = DagBag(process_subdir(args.subdir))
-    dags = dagbag.dags.values()
-    print(_tabulate_dags(dags, tablefmt=args.output))
+    AirflowConsole().print_as(
+        data=sorted(dagbag.dags.values(), key=lambda d: d.dag_id),
+        output=args.output,
+        mapper=lambda x: {
+            "dag_id": x.dag_id,
+            "filepath": x.filepath,
+            "owner": x.owner,
+        },
+    )
 
 
 @cli_utils.action_logging
+@suppress_logs_and_warning()
 def dag_report(args):
     """Displays dagbag stats at the command line"""
     dagbag = DagBag(process_subdir(args.subdir))
-    print(tabulate(dagbag.dagbag_stats, headers="keys", tablefmt=args.output))
+    AirflowConsole().print_as(
+        data=dagbag.dagbag_stats,
+        output=args.output,
+        mapper=lambda x: {
+            "file": x.file,
+            "duration": x.duration,
+            "dag_num": x.dag_num,
+            "task_num": x.task_num,
+            "dags": sorted(ast.literal_eval(x.dags)),
+        },
+    )
 
 
 @cli_utils.action_logging
+@suppress_logs_and_warning()
 def dag_list_jobs(args, dag=None):
     """Lists latest n jobs"""
     queries = []
@@ -324,6 +322,7 @@ def dag_list_jobs(args, dag=None):
     if args.state:
         queries.append(BaseJob.state == args.state)
 
+    fields = ['dag_id', 'state', 'job_type', 'start_date', 'end_date']
     with create_session() as session:
         all_jobs = (
             session.query(BaseJob)
@@ -332,15 +331,16 @@ def dag_list_jobs(args, dag=None):
             .limit(args.limit)
             .all()
         )
-        fields = ['dag_id', 'state', 'job_type', 'start_date', 'end_date']
-        all_jobs = [[job.__getattribute__(field) for field in fields] for job in all_jobs]
-        msg = tabulate(
-            all_jobs, [field.capitalize().replace('_', ' ') for field in fields], tablefmt=args.output
-        )
-        print(msg)
+        all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs]
+
+    AirflowConsole().print_as(
+        data=all_jobs,
+        output=args.output,
+    )
 
 
 @cli_utils.action_logging
+@suppress_logs_and_warning()
 def dag_list_dag_runs(args, dag=None):
     """Lists dag runs for a given DAG"""
     if dag:
@@ -361,13 +361,19 @@ def dag_list_dag_runs(args, dag=None):
         execution_end_date=args.end_date,
     )
 
-    if not dag_runs:
-        print(f'No dag runs for {args.dag_id}')
-        return
-
     dag_runs.sort(key=lambda x: x.execution_date, reverse=True)
-    table = _tabulate_dag_runs(dag_runs, tablefmt=args.output)
-    print(table)
+    AirflowConsole().print_as(
+        data=dag_runs,
+        output=args.output,
+        mapper=lambda dr: {
+            "dag_id": dr.dag_id,
+            "run_id": dr.run_id,
+            "state": dr.state,
+            "execution_date": dr.execution_date.isoformat(),
+            "start_date": dr.start_date.isoformat() if dr.start_date else '',
+            "end_date": dr.end_date.isoformat() if dr.end_date else '',
+        },
+    )
 
 
 @provide_session